Train-Attention: Meta-Learning Where to Focus in Continual Knowledge Learning

Authors: yeongbin  seo
dongha  lee
jinyoung  yeo
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 08:48:14

Abstract: Previous studies on continual knowledge learning (CKL) in large language models (LLMs) have predominantly focused on approaches such as regularization, architectural modifications, and rehearsal techniques to mitigate catastrophic forgetting. However, these methods naively inherit the inefficiencies of standard training procedures, indiscriminately applying uniform weight across all tokens, which can lead to unnecessary parameter updates and increased forgetting. To address these shortcomings, we propose a novel CKL approach termed Train-Attention-Augmented Language Model (TAALM), which enhances learning efficiency by dynamically predicting and applying weights to tokens based on their usefulness. This method employs a meta-learning framework that optimizes token importance predictions, facilitating targeted knowledge updates and minimizing forgetting. Also, we observe that existing benchmarks do not clearly exhibit the trade-off between learning and retaining, therefore we propose a new benchmark, LAMA-ckl, to address this issue. Through experiments conducted on both newly introduced and established CKL benchmarks, TAALM proves the state-of-the-art performance upon the baselines, and also shows synergistic compatibility when integrated with previous CKL approaches. The code and the dataset are available online.

Main paper: 
 introduction: 
 Large language models (LLMs), pre-trained on extensive text corpora, have demonstrated remarkable effectiveness when fine-tuned or prompted to perform a variety of downstream tasks [Brown et al., 2020 , Raffel et al., 2020 , Sanh et al., 2021 , Wei et al., 2021] . However, as the world changes and new knowledge needs to be updated to the parameters, these models often suffer from a significant loss of previously learned knowledge (i.e., catastrophic forgetting [Kemker et al., 2018 , Kirkpatrick et al., 2017] ). To address this issue, the field of continual knowledge learning (CKL) is being actively researched [Jang et al., 2021 [Jang et al., , 2022]] , which aims to teach a model new knowledge while minimizing forgetting of previous knowledge. Previously explored approaches are broadly categorized into three:(1) minimizing parameter changes through regularization, (2) training the expanded parameters of the adapter while freezing the base model parameters, and (3) reviewing old knowledge. However, these approaches naively inherit the inefficiency of the standard fine-tuning procedure of causal LMs, which uniformly apply weights to all tokens, regardless of their importance. This inefficiency of uniform weighting becomes more significant within the context of CKL, where the model is assumed to possess a substantial amount of world knowledge and grammatical capabilities already, thus emphasizing the need for limiting targets of learning. For example, consider a causal LM that has undergone both pre-training and fine-tuning and now requires to update the new information that "The president of the US is Biden." Figure 1a illustrates how the model processes the example sentence. The only sequence that carries the essential information of this sentence is the final sequence s 6 ("The president of the US is" → "Biden"), which encapsulates the context of "US", "president", and "Biden". Conversely, another sequence such as s 4 ("The president of the" → "US") only contains information that is already familiar to the model: the close association between "president" and the name of a nation, as well as the grammatical rule that a noun follows "the". Moreover, s 1 ("The" → "president") introduces a harmful bias, suggesting that "president" should invariably follow "The", although any nouns could follow "The". If the model overemphasizes the likelihood of this sequence, several issues can arise: (1) Parameters will be updated more than the necessary amount to learn only essential information, thus resulting in more forgetting. (2) The training steps required to learn the important sequence could become prolonged.Therefore, we hypothesize that focusing learning efforts on important tokens elevates the performance of the CKL. We present empirical evidence of this in §4.1.1 (paragraph of the analysis on Oracle). The concept of selecting important tokens has been previously explored outside the domain of CKL by Hou et al. [2022] , Lin et al. [2024] , and demonstrates enhanced performance on downstream tasks. These methods share the same principle, assigning more importance (we denote this "token importance") to the token with higher classification error, which assumes a definition of token importance as "tokens with low-confidence are important". While this approach can accelerate learning of low-confidence tokens, it is still not guaranteed that such low-confidence tokens are truly "important". This emphasizes a need for a more comprehensive definition of "token importance". To clarify this, in the example of Figure 1 , it is necessary to consider why human intuition easily accepts that the sequence s 6 is more important than others. This understanding comes from the anticipation that knowing the new president will be useful in the future (e.g., conversation with neighbors, school exams, etc) [Land and Furneaux, 1997] . Building on this concept, we define "token importance" as the expected utility of the token in related tasks, a concept we refer to as usefulness. Upon this definition of token importance, we propose a novel approach to CKL, named Train-Attention-Augmented Language Model (TAALM), which predicts weights of each token based on their usefulness, leveraging this weight on the training phase to enable efficient update of new knowledge. Train-Attention, the supportive model that predicts weight for the base model, is trained through the meta-learning method.We also introduce a new CKL benchmark, LAMA-CKL, designed to offer a more clear comparison of learning and retention performance. This benchmark's advantages over the previous standard are explained in §4.3. We experiment on LAMA-CKL and previous CKL benchmark (TemporalWiki [Jang et al., 2022] ), and our method achieves remarkable state-of-the-art performance on both. Our method is compatible with other approaches, and shows enhanced performance when integrated, indicating a synergistic effect. We also compared RHO-1 [Lin et al., 2024] , which is the recent concurrent work on the token selecting method, where ours shows superior performance on CKL benchmarks. Our main contribution can be summarized in three. (1) We propose a novel token weighting approach to the CKL task, with a novel problem definition and meta-learning method.(2) A new benchmark for CKL based on the LAMA dataset. (3) Through extensive experiments, TAALM proves notable improvements over the baselines. 
related works: 
 Continual Knowledge Learning Continual Knowledge Learning (CKL) [Jang et al., 2021] is one variation of Continual Learning (CL), specified to LLM. It is more focused on updating new knowledge without catastrophic forgetting [Kirkpatrick et al., 2017 , Kemker et al., 2018] of previously learned and preservable knowledge. Previous approaches for CL and CKL can be mainly categorized into three: regularization, architectural, and rehearsal. We analyze that the three approaches share a common goal; to minimize changes in parameters from initial points. (1) Regularization: directly controlling the extent of change in the parameters through weight regulation such as L2 [Kirkpatrick et al., 2017 , Zenke et al., 2017 , Lopez-Paz and Ranzato, 2017 , Aljundi et al., 2018 , Chen et al., 2020] .(2) Architectural: freezing the base model parameters and expanding learnable parameters with adapters such as Lora [Houlsby et al., 2019 , Hu et al., 2021 , Wang et al., 2020 , Dettmers et al., 2024] , thereby keeping initial parameters untouched. (3) Rehearsal: method of continually reviewing the data that is employed to train the initial model, ultimately returning the parameters to the initial points [Shin et al., 2017 , Sun et al., 2019 , He et al., 2019 , Rolnick et al., 2019] . In this view, our method is another approach to achieve the same goal, minimizing change of parameters, by filtering objective tokens.Meta-Learning Meta-learning [Finn et al., 2017 , Hospedales et al., 2021] is most commonly understood as learning-to-learn; the process of improving a learning episode, over multiple outer learning episodes. During meta-learning, an outer (i.e., meta) learner is fitted to improve the learning of the inner (i.e., base) model. The meta-learner could be an initial parameter of the base model [Finn et al., 2017] , an optimizer of the base model [Andrychowicz et al., 2016] , or a hyper-parameter of the base model such as learning-rate [Li et al., 2017 , Franceschi et al., 2018] . In this view, our meta-learner (Train-Attention) is an LLM architectural model that predicts hyper-parameters of the base model, as the token weights serve as the hyper-parameters in the training objective.Token Selecting Methods to enhance learning by selecting specific tokens have previously been explored through various approaches: Token-Dropping [Hou et al., 2022] , Focal Loss [Lin et al., 2017] , and RHO-1 [Lin et al., 2024] . These methods share a common principle: assigning more importance to the token with higher classification error.3 Train-Attention-Augmented Language Model (TAALM) 3.1 Token Importance and Token-Weighted Learning (TWL)PPL θ = -1 N i log p(x i |x <i ; θ) (1) = -1 i w0 i log p(x i |x <i ; θ) × w 0 (w 0 = 1) (2) Token Weighted (TW) PPL θ = -1 i wi i log p(x i |x <i ; θ) × w i (3)To learn a document data D = {x 1 , ..., x N } which is defined as a series of tokens, a dominant causal language model (LM) (θ) commonly employs perplexity (PPL) as the objective function, as formalized in Eq.( 1). This can be also interpreted in the form of Eq.( 2), which assigns a uniform weight (w 0 = 1) to log probabilities of each sequence x i |x <i across all documents. In contrast, our proposed methodology assigns weights 0 < w i ≤ 1 to log probabilities of each sequence, which approximates the importance of each sequence, named token importance [Hou et al., 2022] . We denote this set of weights as the token weight, and the training process that incorporates the weights (Eq.( 3)) as Token Weighted Learning (TWL). 
train-attention: meta-learning to predict token importance: 
 We suggest defining token importance as usefulness, which indicates how much the contained information is useful for solving related tasks in the future. Under this definition, a meta-learning approach can be derived to develop a supportive model (meta-learner) that predicts the optimal token weights. Let θ represent a base causal LM that continually learns knowledge and solves tasks. T D represents a task that can be solved using information contained in D. Training dataset D is a set of pairs of D and T D . We assume a task T D can be defined as any type (e.g., predicting object labels, classification) as long as the performance can be measured in a differentiable form. The set of token weights, denoted as W D , comprises weights w i that represent the importance of each sequence x i |x <i within D. The meta-learner, named Train-Attention and denoted as ϕ, predicts W D from D.As illustrated in Figure 2a , ϕ inherits the architecture and pretrained parameters of the causal LM, but the decoder layer is adjusted to yield only a single-dimensional float between [0,1] for each position.The desired process, learning knowledge and solving a task, is described in two steps; (a) learn: θ is trained on D and is updated to θ ′ . This update occurs in a TWL manner, with a token weight W D,ϕ ← ϕ(D) that ϕ predicts upon observing the data D. (b) solve: The revised model θ ′ is applied to solve the task T D , and the loss value L θ ′ (T D ) is computed to quantify the performance on T D , where L stands for loss function.Two steps can be regarded as one black box function, which receives ϕ as an input and outputs L θ ′ (T D ). In other words, the task performance of θ ′ depends on how ϕ gives attention when learning evidence text data. And ϕ can be optimized to minimize the L θ ′ (T D ). For this, the procedure of (a) and (b) is developed to corresponding steps of Eq.( 4) and ( 5), where α and β are the learning rates for each respective update. First, base model θ is updated to θ ′ through TWL, with the token weight W D which is generated from ϕ. Second, the meta-learner ϕ is updated based on the task performance, L θ ′ (T D ). These two steps of update can be also interpreted like Figure 3 . As the model θ steps out to a new state θ ′ , the resulting position depends on which token weight (W ) is applied. Some positions are closer to the θ * , a model optimally capable of the task T D , as the distance is measured with L θ ′ (T D ). We can conclude the weight with a shorter distance (W * D ) is more optimal, in the perspective of usefulness. To prevent the θ from converging to the point θ * , which disables the measurement of distances, we reset the model parameters to the initial state θ after every update of ϕ. More detail is depicted in Figure 4 and Algorithm 1. As the gradients of parameters of ϕ are tracked during the updating of θ, its actual implementation is akin to the second derivative. The max iteration step of θ (denote as M in Algorithm 1) is fixed to 1 through our experiment. We employ gradient accumulation when updating ϕ for batch effect.EQUATIONL learn = TWPPL θ (D, W D,ϕ ) Evaluate ∇ θ L learn Update: θ ← θ -α∇ θ L learn end for (θ updated to θ ′ ) L solve = L θ ′ (T D ) Evaluate ∇ ϕ L solve Update: ϕ ← ϕ -β∇ ϕ L solve Reset θ ′ to initial point θ end whileOn the inference phase, θ learns data in TWL manner as in Eq.( 4), with the parameter of ϕ frozen. This system is Train-Attention-Augmented Language Model (TAALM). In this work, the foundational structure of ϕ is fixed to the small model (TinyLlama-1.1B [Zhang et al., 2024] ), while it is augmented to both large (Llama2-7B [Touvron et al., 2023] ) and small base models, because ϕ is compatible with any base model that shares the same tokenizer. Additionally, we explore utilizing a 101M-parameter bidirectional transformer (BERT) [Devlin, 2018] as a Train-Attention (TA) to further reduce resource requirements. 
experiment: 
 We conduct experiments on two benchmarks. One is our newly designed LAMA-CKL, and the other is the established benchmark, TEMPORALWIKI [Jang et al., 2022] . We exclude the CKL benchmark by Jang et al. [2021] which is not publicly available. In this section, we present the corpus, evaluation setup, and training detail for Train-Attention and the test result within our proposed LAMA-CKL benchmark. For TEMPORALWIKI, most configurations are aligned with the original work. For LAMA-CKL, we tailor the LAMA dataset (LAnguage Model Analysis [Petroni et al., 2019] ) to assess the CKL performance, especially the T-REx [Elsahar et al., 2018] part which consists of data from Wikipedia and Wikidata. LAMA is a cluster of datasets that measures how much world knowledge is contained in the LLM. Each unit in the dataset includes a knowledge base triple <subject, relation, object>, along with corresponding evidence documents that support the information contained in this triple. Referring to the previous work [Jang et al., 2022] , a CKL benchmark should evaluate both plasticity and stability. Plasticity refers to how well the model updates new knowledge, while stability refers to how little the model forgets existing knowledge. Accordingly, we sample 500 units of TO-LEARN and NOT-TO-FORGET sets from LAMA to assess each dimension. During the evaluation, as illustrated on Figure 5 , the model learns the evidence documents in the TO-LEARN set. It is then tested on both the TO-LEARN task and the NOT-TO-FORGET task to assess plasticity and stability, respectively.Dataset Setup Here, we outline a protocol for sampling test corpora for the LAMA-CKL benchmark. As TO-LEARN set represents "the knowledge that the model either encounters for the first time or needs to update", it is selected based on two constraints: (1) sample from the categories of time-variant relations, predicated on the assumption that knowledge categorized as time-variant typically requires updates. (2) to ensure the concept of "knowledge new to the model", we select units where the task accuracy is zero when measured with pre-update baselines. Conversely, because NOT-TO-FORGET set represents "the knowledge that the model already knows and aims to retain", it is selected from categories of time-invariant relations, with task accuracy of 1. We recommend sampling a new dataset by the specified constraints when evaluating models outside of the LLaMA-family, for more accurate assessment. The categorization of time-variant and time-invariant follows Jang et al. [2021] . Each selected unit includes (1) an evidence document, (2) a knowledge base triple (e.g., <Lochinvar, is an instance of, castle>), and (3) a descriptive sentence encapsulating the triple (e.g., "Lochinvar is a castle"), which is inherited from LAMA dataset. The task is predicting object label tokens in the descriptive sentence. Details on data are in Appendix A.1Evaluation Setup During the evaluation, each epoch consists of both a training phase and a test phase. In the training phase, the model is trained on a set of 500 evidence documents from the TO-LEARN set. Subsequently, in the test phase, the model's prediction accuracy for the object labels is assessed using 500 descriptive sentences from both the TO-LEARN and NOT-TO-FORGET sets. This process is repeated over 30 epochs. For the TO-LEARN set, an increase in mean accuracy from 0 signifies the model's plasticity. Conversely, a decline in mean accuracy for the NOT-TO-FORGET set from 1 to lower values indicates the model's stability, as it tends to forget previously learned information.In the proposed benchmark LAMA-CKL, we suggest four main factors as evaluation indicators. 1) Top Acc: the highest TO-LEARN accuracy among checkpoints of 30 epoch. 2) Epoch: the epoch where the Top Acc appears. 3) NF Acc: NOT-TO-FORGET accuracy of the checkpoint model which is the same as Top Acc. 4) Total Knowledge: the sum of Top Acc and NF Acc, indicating total capacity of knowledge including updating and maintaining. We chose these factors because the best CKL system is one that learns the most and the fastest and loses the least. Factor 1, 3, and 4 are better if higher, while factor 2 is better if lower. The detailed configurations for training datasets are in Appendix A.1.1. Baseline Setup We utilize Llama2-7B integrated with QLoRA [Dettmers et al., 2024] as a base model. The baseline methods and their hyper-parameter settings follow previous CKL study of Jang et al. [2022] : standard finetune, K-Adapter [Wang et al., 2020] , Mix-review [He et al., 2019] , LoRA [Hu et al., 2021] , RecAdam [Chen et al., 2020] . We regard standard finetune on QLoRA as a substitute for full finetuning and LoRA, thus skipping the two baselines. We also compare RHO-1 [Lin et al., 2024] , which is the most recent concurrent work on the token selecting method, sharing a similar concept with ours. We chose the initial parameter state as a reference model, which is utilized to select important tokens for RHO-1, with other hyper-parameters following the optimal of the original. We also evaluate a model trained in TWL manner with Oracle token weight. For which, a weight of 1 is exclusively assigned to the object label token in the evidence document, and the rest is assigned zero weight. Oracle is compared for two purposes: 7 and Table 1 , our method (TAALM) overwhelms other baselines on ability and speed of learning. Ours records the Top Acc of TO-LEARN task as 0.4290 on only 4 epochs of updating. The accuracy record of ours is 3.04 times higher than the second place (RHO-1). And the required epoch is only 25% of the second place (standard finetune).K-Adapter shows the highest NF Acc, rarely forgetting previous knowledge. However, the TO-LEARN accuracy also barely increases, indicating that parameter updates rarely occur. This could be due to the architectural difference from QLoRA, the backbone of other baselines. Therefore, we conclude that it is meaningless to compare K-Adapter and QLoRA based baselines in the same condition.When comparing only the QLoRA base methods, TAALM shows overwhelming performance in all dimensions. Ours learn the most, the fastest, and forget the least. RHO-1, one of the token selecting methods, also shows a bigger capacity for learning and less forgetting than standard finetuning, However, the advantage is minimal compared to ours.Combination with TAALM improves all of the previous baselines As Train-Attention is an approach to manipulating loss values on the end side, it is easily compatible with previous methods based on other concepts. Thus we combine various methods to TAALM and observe their performance.Referring to the experimental results and details in Appendix A.4.2, each combined version shows a highly improved capacity than the baseline alone. Especially, combining our method with K-Adapter demonstrates a considerable balance between stability and plasticity.Oracle vs Train-Attention vs Standard finetune It is assumed that TWL with Oracle label is the upper bound of performance of TWL with Train-Attention, because it is sharply targeting only TA on BERT achieves comparable performance with small resources We train BERT as a TA and evaluated it on the LAMA-CKL dataset. Training TA on BERT requires only a single 24GB GPU, significantly reducing resource usage compared to the previous model (single 82GB GPU), yet achieving performance similar to the larger (1.1B) TA (Appendix D).We conduct an ablation study on various design choices applied to the token importance predicted by TA: (1) masking out tokens (setting importance to 0) in real-time when the prediction matches the label, and (2) dropping weights with token importance below the top k% threshold. The ablation study reveals that heuristic adjustments degrade performance, as TA is already in an optimized state. Details are in Appendix E.We experiment on the original CKL learning benchmark TEMPORALWIKI [Jang et al., 2022] , where models have to continually learn Wikipedia documents of serial periods (0809, 0910, 1011, 1112) and test on the corresponding TWIKI-PROBES, which is a dataset of knowledge base triples. As we train Train-Attention on the 0809 data, tests are conducted on the rest. We experiment with only a small (TinyLlama-1B) model, which is bigger than the baselines of the original work (GPT-2 Large). We conduct a separate experiment on QLoRA-based K-Adapter based models, referring to the analysis on experiment of LAMA-CKL on §4.1.1. We only consider training DIFFSET, which is the only changed part of Wikipedia, because it is reported as a condition of the best performance. Most of the experimental settings follow the original, and additional change is described in Appendix C.Referring to Table 2 , our method presents the state-of-the-art performance across both experiments on QLoRA based and K-Adapter based models. This achievement is consistent in all periods, and in both CHANGED and UNCHANGED TWIKI-PROBES. This result is aligns with the LAMA-CKL benchmark result, showing that our method has a substantial advantage on the CKL. QLoRA based baselines showed poor performance compared to K-Adapter based baselines, indicating architectural disadvantage. We also test TAALM optimized for LAMA-CKL on the TEMPORALWIKI, referring For the benchmark TEMPORALWIKI, because DIFFSET is corpora of evidence documents for CHANGED set, learning of DIFFSETS is supposed to result in performance improvement over CHANGED set and forgetting of UNCHANGED set. However, during our experiments, we observe that both CHANGED and UNCHANGED performance tend to move in similar directions when learning DIFFSET, which is in contradiction to our assumption (Appendix C.4). We analyze this for two primary reasons. First, the DIFFSET contains evidence documents for both the CHANGED and the UNCHANGED sets (Appendix C.4). This is a complicating factor in the evaluation of stability. Second, the experimental setup involves training on a vast amount of data, an average of 707K documents per period, for just a single epoch at a low learning rate. This might result in learning little amount of knowledge, while the task ability is challenged by extensive iterations of updates; which is closer to a continual learning setup rather than CKL. To address this issue, we structured the LAMA-CKL as follows:(1) To mitigate the issue of data overlap, we partition the dataset into variant and invariant subsets. These subsets are further classified based on the task accuracy measured by pre-update baselines.(2) We conduct training over multiple epochs on a relatively small dataset to observe the acquirement of knowledge. In practice, our benchmark shows a clear upward trend in the TO-LEARN set and a distinct decline in the NOT-TO-FORGET set as training progresses, clearly demonstrating the contrast between plasticity and stability. 
conclusion and limitation: 
 In this paper, we demonstrate that the application of Train-Attention significantly enhances CKL performance and is also synergistic with other baselines. Nevertheless, our work has the following limitations and potential for future exploration.Task specificity of Train-Attention Train-Attention is trained to focus on information related to tasks encountered in the training session. This allows task performance to increase, but on the other hand, it left questions as to whether it would be possible to cope with other tasks. Nonetheless, if the task entails the acquisition of general knowledge, it will be transferable to other tasks sharing similar distributions. For instance, TAALM optimized to LAMA-CKL also achieved the best performance on the TEMPORALWIKI ( §4.2.1). Additionally, Train-Attention can ever evolve to adapt, enabling optimal performance for the current tasks.What if there are no data-task pair Train-Attention can be trained only if there is a data-task pair.If there is no paired dataset, it can be doubted that training is difficult. However, every knowledge has its purpose, and we can find a strategy to discover it. 1) Search: When the data and task pools are separate, we can join highly related pairs via searching. TEMPORALWIKI is also a dataset in which data and tasks are not paired, thus we conduct a lexical search. In the future, also dense research methods can be explored. 2) Generate: If there is even no separate task pool, we can at least get prior information about what kind of tasks are probable to come in the future. Synthetic tasks can be generated via instruction-tuned LLM, based on this prior information. These methods also resemble the human's cognitive strategy, who often revisit past memories and pose hypothetical questions to themselves to enhance the efficiency of their learning processes.Broader impacts Our method aims to increase the ability of CKL, therefore enhancing the practicability of LLMs and saving the computational resources for fine-tuning entire huge LLMs. We believe that this paper does not have any immediate negative societal impact.F  
 
 Appendix: As each unit of the LAMA dataset contains multiple evidence documents, we sample one document for each unit, based on the following criteria: (1) The document should exceed a length of 70 tokens.(2) The document must include both the subject and object entities. These documents are then truncated to a maximum length of 512 tokens. We assess the accuracy of the knowledge base triple of each unit to filter them. We utilize the TinyLlama-1.1B Zhang et al. 4 . Table 3 presents data statistics for NOT-TO-FORGET and TO-LEARN sets. 5 describe performances of small (TinyLlama-1B) baselines on the LAMA-CKL.The performance of TAALM based on the small (TinyLlama-1B) architecture is significantly better even when compared to the large (Llama2-7B) baselines. Referring to Table 6 and Figure 12 , the combination of ours and Mix-review shows the highest Top Acc and Total Knowledge among all variations. Unlike Mix-review alone, the Mix-review combined with ours applies Train-Attention not only to the train data but also to the review data. As a result, it is observed that the training becomes more stable than the original, which tends to be unstable due to the doubled amount of data (train + review). This indicates that the combination of Train-Attention with the rehearsal approach is synergistic. Variations of "ours + K-Adapter" show relatively low learning capability (Top Acc, Total Knowledge) while featuring still robust stability. Therefore, a combination of ours and K-Adapter can be considered a balanced variation between learning and maintaining.Figure 13 presents the trade-off between plasticity (TO-LEARN) and stability (NOT-TO-FORGET) of all baselines, including all variations of combination. Variations combined with our method are exclusively positioned in the upper right quadrant, indicating a minimal trade-off between updating and forgetting. In the main experiment (Figure 7 ), the accuracy declines after a rapid peak after about the fourth epoch. We hypothesize this is a result of overfitting. While the token weights from TA include beneficial targets, there must also be some that are not. Performance achieves peak until the model completes learning for the true target, but learning may continue for the false targets. This continued learning leads to parameter updates in suboptimal directions, resulting in forgetting.The comparison between TA and the oracle (Figure 17 ) provides evidence for this rationale. As the oracle's performance does not decline with further training, the differences in accuracy trends could be caused by false targets. This phenomenon can also be interpreted as a type of overfitting where the model becomes too fitted to the training data which has different distribution from test data. As the phenomenon of declining after peaking is common due to overfitting, this cycle seems to occur more rapidly for TAALM. Each unit of the LAMA dataset comprises a knowledge base triple and a corresponding sentence that encapsulates the information contained within the triple. These sentences are presented in a short descriptive form. In contrast, our newly proposed schematic form organizes these triples more systematically. The template and example are described in Figure 15 . We propose the schematic form because it better aligns with the uni-directional nature of causal language models (LMs). In the descriptive form, critical information often comes after the label tokens, which can be problematic. For instance, in the descriptive form template "[X] is [Y] citizen", the causal LM lacks the crucial cue 'citizen' when predicting [Y] . This not only makes the assessment less accurate but also introduces noise into the train-attention learning process. The original study shows that learning only the DIFFSET results in better performance than learning the entire snapshot. So we excluded the baselines that learn full snapshots and set learning of DIFFSET as the default. We also employed heuristic filtering for the DIFFSET. Empty texts containing only 'nan' in the content, or texts with more than 70% of non-letters, which are assumed to contain little meaningful information, are filtered out. Therefore the total number per DIFFSET data decreases from an average of 837K to 707K documents per each dataset.Each Twiki-Probes consists of two parts: (1) CHANGED: assumed to consist only of knowledge base triples that contain changed information compared to the last period.(2) Unchanged: contain only retained information from the last period, which is the complementary set of the CHANGED. CHANGED set is supposed to measure plasticity, while UNCHANGED set is supposed to measure stability.We follow the evaluation settings of the original work Jang et al. [2022] : a model continually learns train datasets of each period while evaluating performance on TWiki-Probes after training of each period. The evaluation metric is the mean perplexity of object label words. One dataset is updated for only 1 epoch.The original study simply arranges the triples in order (e.g., "Wales official language Welsh") and sets the task as predicting the object label token. As this form of sentence might seriously violate the grammar, it has been reported that perplexity tends to pick extraordinarily high in the original work.To handle this, we adopt the schematic task format used in LAMA-CKL ( §4.1).Baselines and hyper-parameters settings follow that of LAMA-CKL ( §4.1), which follow the Jang et al. [2022] .The DIFFSET and Twiki-Prob of the first period (0809) are used for training Train-Attention. To train Train-Attention (Algorithm 1), data (D) and the corresponding task (T D ) must be paired, while the documents in DIFFSET and knowledge base triples in Twiki-Probes are not paired. To tackle this challenge, we search the DIFFSET documents which contain both the object and subject of the Twiki-Probes unit and manually join them. The total number of pairs is 5235. Validation setting and hyper-parameters follow §4.1.In this section, we present the detailed experimental result on TEMPORALWIKI, specifically involving the discussion in §4.3, on the reason we propose the new benchmark LAMA-CKL. Figure 16 indicates that performances in both Changed and Unchanged mostly move in the same direction when the baselines are trained on the DIFFSET of each period, presenting a rare trade-off between plasticity (Changed) and stability (Unchanged). This is contrasting to the observation on our new benchmark, LAMA-CKL (Figure 13 ), which exhibits a clear trade-off between plasticity (TO-LEARN) and stability (NOT-TO-FORGET). This suggests that the LAMA-CKL is a better benchmark for observing the trade-off between plasticity and stability.Table 7 presents the portion of task sentences in each Changed and Unchanged set, which have evidence documents in each period of the training dataset (DIFFSET). An average of 8.0% of every Unchanged set have evidence documents in the DIFFSET of the corresponding period, indicating that training on DIFFSET partially results in an improvement in the performance on the UNCHANGED set. The Unchanged set from the last period (TWiki-Probes-1112) contains evidence documents in the DIFFSETs from all three preceding periods, suggesting that its supportive effect is cumulative. This is a complication factor for observing stability. Due to the substantial GPU resources required to train the TA, it is necessary to find ways to reduce resource consumption. A promising approach is utilizing Bidirectional Transformer (BERT) as a body for TA, which has high inferential capabilities even at a very small size (108M) compared to the previous body (Tinyllama 1.1B), due to its bidirectional property. Since BERT has a different tokenizer from our generation model, the Llama family, we integrate BERT with the Llama2 tokenizer and pre-train it for one epoch on 17GB Wikipedia documents (9 days using 8 of 24GB GPUs). Then, we finetune this BERT as TA, paired with the generation model of 1B (Tinyllama). This very lightweight TAALM is sufficiently trained on only a single 24GB GPU, significantly reducing resource use compared to the previous version (single 82GB GPU), thus making it affordable for the general environment. On the inference, the TA on BERT demonstrates compatibility with both the 1B and 7B generation models. Although its performance is below that of the TA on Llama, it still exhibits the highest performance among the other baselines. Based on the token importance predicted by the TA, various design choices are possible. We explore and compare the effectiveness of these variations. This study enhances our understanding of how generative LM interacts with token weights when learning data. The description of the components and experimental results are as follows.(1) Token-importance weight (ours) : The original variation that utilizes token-importance weight predicted by TA for target-weighted learning.(2) Known token masking : Masking out the tokens in real-time when prediction and label matches. This method is intended to enhance "model awareness" in TA, as TA is more oriented to "task awareness." (3) Token weight dropping : Among token-weight generated by TA, dropping weights that are below the top k% levels. We tested 50% and 80%. Vanilla TA is the same as the threshold of 0%. This method is intended to cut out noisy targets, as TA is supposed to assign lower weight to un-useful tokens.Results Known token masking does not yield better results compared to TAALM w/ tokenimportance weight. We hypothesize that the effect of known masking is limited because task awareness is already achieved when the loss of learned tokens is reduced. Test results on the TAALM w/ token weight dropping show that as the threshold increases, the top accuracy decreases. This suggests that some useful targets are mixed in among the lower weights, and it helps the model learn better somehow. On the contrary, not-to-forget accuracy slightly improves as the threshold increases. This seems as the effect of 1) cutting out noisy targets and 2) trade-off for lower learning. However, Total Knowledge is best on the TAALM w/ token-importance weight (ours). Overall experimental results indicate that, since TA is optimized to maximize task performance, adding heuristic interventions appears to produce suboptimal outcomes. F Analysis on the attention pattern of TA TA is observed to generally assign attention to proper nouns, nouns, and verbs that contain the subject's character. The focus of attention seems to be diverse depending on the content of the text. For autobiographical texts, TA shows a tendency to focus on words that represent the person's occupation or major events (Figure 18a ). In passages listing regional relations, TA pinpoints the names of locations (Figure 18b ). This appears to be due to the consideration of probable queries.While TA (trained on LAMA-CKL) omits some words in the documents, it tends to not miss location names. This is likely because many queries in the LAMA-CKL benchmark involve location-related aspects (e.g., birthplaces, location of the workplace).We also provide an attention map of TA trained on the multi-session chat (Figure 18c ). Here, we regard prior dialogue sessions as data (D) and an understanding of the next session as task (T D ). Unlike Wikipedia documents, chit-chat dialogues contain fewer useful words, highlighting the necessity for TA. TA focuses on the interlocutor's information like the occupation and pet's name. https://github.com/ybseo-ac/TAALMIn this work, we use the notation xi|x<i to denote that the element xi is sequenced after the preceding elements x0, x1, ..., xi-1.Directly referred from 'Relations of InvariantLAMA ofJang et al. [2021] 


Decision: This paper investigates how to improve training efficiency of continual knowledge learning in large language models. It presents a meta-learning approach to estimating the importance/weights of tokens, reducing unnecessary parameter updating, and mitigating catastrophic forgetting. It further proposes a new benchmark for evaluating the trade-off between learning and retaining. Experiments demonstrate the effectiveness of the proposed method.

Code and datasets are provided for reproducing the experimental results.

All reviewers found this paper provides novel approach and new benchmark for continual knowledge learning which would be significant contributions to the ML community. After the rebuttal and discussion, all reviewers keep their positive scores. I recommend acceptance.

Reviews: review: paper_summary: paper_summary: The paper introduces Train-Attention-Augmented Language Model (TAALM), a novel approach to continual knowledge learning (CKL) in large language models (LLMs). Unlike traditional methods that uniformly apply weight across all tokens, TAALM dynamically predicts and applies weights to tokens based on their importance using a meta-learning framework. This approach aims to enhance learning efficiency by targeting essential knowledge updates and minimizing forgetting. The authors also introduce a new benchmark, LAMA-CKL, to better assess the trade-off between learning new information and retaining existing knowledge. Experimental results show that TAALM significantly outperforms existing CKL methods on both new and established benchmarks.
strength_weakness: strength_weakness: Generalization to Other Tasks: While the paper demonstrates the effectiveness of TAALM in the context of CKL, it would be beneficial to explore its applicability to other types of continual learning tasks beyond language models.
Computational Resources: The training of TAALM, especially with large models, requires significant computational resources. A discussion on the scalability and efficiency of the approach for smaller models or resource-constrained environments would be useful.
Ablation Studies: While the paper includes comprehensive experiments, additional ablation studies to isolate the impact of different components of TAALM (e.g., the specific meta-learning algorithm used) would strengthen the evaluation.
questions: questions: Can the authors provide more details on the potential applicability of TAALM to other types of continual learning tasks outside of language models?
How does TAALM perform with significantly smaller models or in environments with limited computational resources?
Could the authors include more detailed ablation studies to better understand the contribution of each component of TAALM to the overall performance?
limitations: limitations: The paper adequately discusses the limitations of the proposed approach, particularly the task-specific nature of Train-Attention and the requirement for data-task pairs for training. The authors also highlight the potential for Train-Attention to evolve and adapt to new tasks, suggesting areas for future exploration. However, a more detailed discussion on the computational requirements and scalability of TAALM would be beneficial.
paper_summary: paper_summary: The article introduces Train-Attention Augmented Language Model (TAALM), a novel approach for continual knowledge learning in large language models. TAALM dynamically assigns weights to tokens based on their usefulness, optimizing learning efficiency and minimizing forgetting.
strength_weakness: strength_weakness: - The proposed method is not novel and it is widely used in multi-task learning to adjust task loss weights in the learning process, like "MetaWeighting: Learning to Weight Tasks in Multi-Task Learning".
- This paper lacks of explanation of why it could maintain previous knowledge without forgetting.
questions: questions: - What will be the influence of dropping the lightweight token sequences in training? If only important token sequences matter in learning new knowledge.
- Can authors offer more theoretical or experimental explanations for the effectiveness of the trained attention?
- I was confused about the relation of the learned meta weight with continual learning. Why dropping the unuseful tokens will help the model learn without forgetting? 
- The ablation study of the learned meta-weight is missing. I only found the results in Figure 6, however the effectiveness of the learned weight should be further validated.
limitations: limitations: N/A
paper_summary: paper_summary: This paper studies the continual knowledge learning (CKL) problem in large language models. The authors notice that the existing methods either apply equal weights to all tokens or re-weight tokens with a trivial consideration that tokens with low confidence are important. They propose a new definition of token importance based on the expected functionality of the token in solving related tasks. In particular, they design a meta-learning strategy to learn token weights in the training phase to facilitate the continual learning of new knowledge while retaining already learned knowledge. A new benchmark, LAMA-CKL, is established. The experimental results demonstrate that the proposed method outperforms other methods in both existing and introduced benchmarks.
strength_weakness: strength_weakness: 1. Naming token importance as "usefulness" is too broad and does not accurately reflect the motivation of the paper.
2. Meta learning is a well-known method, the description of the algorithm details in Section 3.2 is however confusing and unclear.
3. The proposed method seems to only be applicable to the experimental setup described in the paper. This setup requires pairs of \mathcal{D} and \mathcal{T}_\mathcal{D}, where \mathcal{T}_\mathcal{D} represents a task that can be solved using the information contained in \mathcal{D}. This setup is not general enough and is limited to very specific scenarios.
questions: questions: Why does the proposed meta learning approach reach performance peaks in fewer epochs compared to other methods? Is there an explanation for this?
limitations: limitations: The discussion of limitations is included in the paper.
paper_summary: paper_summary: This paper proposes a meta-learning framework to dynamically adjust token perplexity weights based on their usefulness to achieve good continuous knowledge learning performance and also provides a new benchmark for CKL.
strength_weakness: strength_weakness: 1. The idea of adjusting token weights is not fresh enough and shares some similarity with [1]. However, training a meta-learner to evaluate token importance is still a good try.

2. The analysis about token importance (Figure 6 and its corresponding analysis) is not enough. Analysis about attention pattern on these tokens would be a good supplement.


[1] Lin Z, Gou Z, Gong Y, et al. Rho-1: Not all tokens are what you need[J]. arXiv preprint arXiv:2404.07965, 2024.
questions: questions: 1. In Figure 7, I am wondering why the to-learn accuracy drops after 4-th epoch? It is kind of counter-intuitive to me as I expect it will grow along with increasing epoch.

2. Could the authors provide more analysis and rationale behind the phenomena that K-adapter keeps its not-to-forget accuracy basically unchanged?
limitations: limitations: The authors addressed the limitations adequately.

Comments: title: Re: Official Comment by Reviewer 6R4i
Sincere thanks
Sincere thanks
Sincere thanks
Sincere thanks
Discussion
Respond
Response to Rebuttal  by Authors
comment: Dear Reviewer 6R4i,

Having gone through the paper and the discussion here, it appears that the general consensus is leaning towards acceptance. If the authors have addressed your concerns. Please adjust your scores accordingly.

Regards,

AC
We are glad that our reply addressed your concerns. Your feedback has strengthened our work. We will incorporate your feedback into the camera-ready version. We would be happy to address any further questions/suggestions that might come up until the end of the discussion period.
We thank again for your constructive review. We hope that our responses have adequately addressed your previous concerns, and we would be happy to address any further questions/suggestions that might come up until the end of the discussion period.
Thank you for recognizing our work and for raising your score. Your feedback has strengthened our work. We would be happy to address any further questions/suggestions that might come up until the end of the discussion period.
We appreciate for the positive feedback on our response. Your review has strengthened our work. We will incorporate your feedback into the camera-ready version. We would be happy to address any further questions/suggestions that might come up until the end of the discussion period.
Dear Reviewers,

Thanks for the reviews. The authors have uploaded their responses to your comments, please check if the rebuttal address your concerns and if you have further questions/comments to discuss with the authors. If the authors have addressed your concerns, please adjust your rating accordingly or vice versa.

Regards,

AC
Many thanks to the author for the responses. They address my concerns and I will maintain a positive score.
Many thanks to the authors for the detailed answers. After considering both the other reviews and the rebuttals, I will increase my score.
Thanks for your detailed response to my questions. I am impressed with the results you achieved with both larger and smaller pre-trained models. I would like to keep my score.