Train-Attention: Meta-Learning Where to Focus in Continual Knowledge Learning
Authors: Yeongbin Seo
Dongha Lee
Jinyoung Yeo
Venue: NeurIPS 2024 poster (2024)
Keywords: continual learning
continual knowledge learning
large language models
meta-learning
train-attention
token weight
Primary area: natural_language_processing
TL;DR: enhancing continual knowledge learning performance through meta-learning based token weighted learning method

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.

Train-Attention Meta-Learning Where to Focus in Continual Knowledge Learning
Yeongbin Seo
Dongha Lee"
Jinyoung Yeo
Department of Artificial Intelligence
Yonsei University { (suhcrates, donalee, jinyeo @yonsei. ac.kr
Abstract
Previous studies on continual knowledge learning (CKL) in large language models (LLMs) have predominantly focused on approaches such as regularization, archi- tectural modifications, and rehearsal techniques 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 novel CKI. approach termed Train-AAtennton Augmented Language Model (TAALM), which enhances learning efficiency by dynamically predicting and applying weights to tokens based on their usefulness. This method employs meta-learning framework that optimizes token importance predictions, facilitating targeted knowledge updates and minimizing forgetting. Also, observe that existing benchmarks do not clearly exhibit the trade-off be- tween learning and retaining, therefore we propose a new benchmark, LAMA-CKL. address this issue. Through experiments conducted on both newly introduced and established CKI. benchmarks, TAALM proves the state-of-the-ar performance upon the baselines, and also shows synergistic compatibility when integrated with previous CKL approaches. The code and the dataset will be available online?)
(a)
Trumph
(b)
Update: "The president of the US Biden"
Predict
Which information will useful in the future?
8
importance
weight
The president
unimportant
84
The president president of the US
85
unimportant
0.09
The
86
unimportant
0.03
Train-Attention Model (Ours)
The
president
of the
US
Biden
Important
0.81
20
Ø1
22
Ø3
27
25
26
Figure 1: (a) Learning of Causal LM: The document is decomposed into multiple token sequences
8t - = which aligns with different importance, but uniformly weighted. (b) Train-Attention: Our proposed Train- Attentionn learns to predict weights that approximate importance, to enable targeted continual knowledge updates through label-free meta-learning method.
"oo-corresponding authors fttpps: / / thub. http//ybsoooo///aaa In this work, the notation 2.|3ec denote that the element x; is sequenced after the preceding elements no, 3i, J4-1-
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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. [02T]] Wei et al. 202T]. However, as the world changes and new knowledge needs to be updated to the parameters, these models often suffer from significant
loss of previously learned knowledge (i.e., catastrophic forgetting [Kemker et al.) 2018 Kirkpatrick]
et al. 2017)) D. To address this issue, the field of continual knowledge learning (CKL) is being actively researched [ann et al. 202T 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 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 I illustrates how the model processes the example sentence. The only sequence that carries the essential information of this sentence is the final sequence s6 ("The president of the US is" + "Biden"), ") which encapsulates the context of "US", "president", and "Biden". Conversely, another sequence such 84 ("The president of the" + "US") only contains information that is already familiar to the model: the close association between "president" and the name of nation, well the grammatical rule that noun follows "the". Moreover, 81 ("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 clevates the performance of the CKL. We present empirical evidence of this in $4.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). [in 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 important". While this approach can accelerate
importance "tokens with low-confidence learning of low-confidence tokens, it is still not guaranteed that such low-confidence tokens are truly "important". This emphasizes a need for more comprehensive definition of "token importance". To clarify this, in the example of Figure I] it is necessary to consider why human intuition easily accepts that the sequence Sg is more important than others. This understanding comes from the anticipation that knowing the new president will be useful in the future (c.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 novel approach to CKL, named Train-Attention-- Ateetiio-AAumennte 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. trained through the meta-lcarning method.
the supportive model that predicts weight for the base model, is
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-lcarning method.

---

(2) A new benchmark for CKL based on the LAMA dataset. (3) Through extensive experiments, TAALM proves notable improvements over the baselines.
2
Related Works
Continual Knowledge Learning Continual Knowledge Learning (CKL.) (Jang al. 2021 is one variation of Continual Learning (CI.), specified to L.M.. It is more focused on updating new knowledge without catastrophie forgetting (Kirkpatrick al.] 2017 Kemker 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 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. [017 Zenke et al. 2017 Lopez-Paz and Ranzato 2017 Aljundi et al. 2018 Chen et al. 2020]
(2) Archifectural: freezing the base model parameters and expanding learnable parameters with adapters such as Lora (Houlsby et al. 2019
20199111 (uu et al. 2021 Wang et al. [2020] Dettmers et al. (2224) 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] [un et al.) [019 He al.) (2199 Rolnick et al., 2019]. this view, our method another approach to achieve the same goal, minimizing change of parameters, by filtering objective
tokens.
Meta-Learning Meta-learning [Finn et al. 2017 Hospedales al. 2021 is most commonly understood as Searning-to-leann the process of improving a learning episode, over multiple outer learning episodes. During meta-learning, of the inner (i.e., base) model. The meta-learner could be an initial parameter of the base model 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-parameten of the base model, as the token weights serve as the hyper-parameter in the training objective.
outer (i.e., meta) learner is fitted to improve the learning
[Finn al.] 2017]. optimizer of the base model [Andrychowicz al. [016],
Token Selecting Methods to enhance learning by selecting specific tokens have previously been explored through various approaches: Token-Dropping [Hou et al.) [022]] Focal Loss [Lin et al. 2017], and RHO-1 [Lin et al. 2024]. These methods share common principle: assigning more importance to the token with higher classification error.
3 Trainn-Atention--uugeee Language Model (TAALM)
3.1 Token Importance and Token-Weighted Learning (TWL)
PPL., - E, logg(osslseei
(1)
= E,
1 Σ, log p [r 0) x wp (wo = 1)
(2)
Token Weighted (TW) PPL. E, Σ, log p(,; [ 0) x wi
To learn a document data D {ai,..., EN} which is defined
series of tokens, a dominant causal language model (LM) (#) commonly employs perplexity (PPL) as the objective function, as formalized in Eq. . This can be also interpreted in the form of Eq. which assigns a uniform weight (wo = 1) to log probabilities of each sequence æi|are across all documents. In contrast, our proposed methodology assigns weights 0 < wy < to log probabilities of each sequence, which approximates the importance of each sequence, named token importance [Hou et al.) [002]. We denote this set of weights as the token weight, and the training process that incorporates the weights (3) as Token Weighted Learning (TWL).
3.2 Train-Attention: Meta-I.crrning 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, meta-learning
3

---

(a) Architecture of Train-Attention
(b) Train-Attentoon-uummenneee Language Model
Replsced
Train-Attention
We
LUM archiacture
LM
LLM arestecturo
TV
Text
Base Model
hare
dlg
uce
Inpu:
(Causal LM)
Infer
PPL
Causal LM
Train-Attention
TWPPL
PPL & Weight
Figure 2: (a) depicts the architecture of Train-Attentionn which shares the structure of causal LM, while the decoder layer (LM head) of causal LM replaced from a linear layer of [hidden size x vocab size] dimension to [hidden size 1] dimension, which is TA (TrainAAteentin head. (b)
depicts the TAALM, where the Train-Attention (Φ) is augmented to the base model (0).
approach can be derived to develop supportive model (meta-leamer) that predicts the optimal token weights. Let represent base causal LM that continually learns knowledge and solves tasks. TD represents task that can be solved using information contained in D. Training dataset L is a set of pairs of 'D and Tp. We assume a task Tn can be defined as any type (e.g., predicting object labels, classification) as long as the performance can be measured in differentiable form. The set of token weights, denoted as WD. comprises weights w; that represent the importance of each sequence r, [ within D. The meta-learner, named Train-Attention and denoted o, predicts WD from D) As illustrated in Figure[2, & inherits the architecture and pretrained parameters of the causal LM, but the decoder layer is adjusted to yield only single-dimensional float between [0,1] for each position.
The desired process, learning knowledge and solving task, is described two steps; (a) learn: is trained on D and is updated to 0'. This update occurs in TWL manner, with token weight WD.s f ((D) that @ predicts upon observing the data D. (b) solve: The revised model O' is applied to solve the task TD. and the loss value Lgr (Tp) is computed to quantify the performance Tp. where L stands for loss function.
Two steps can be regarded as one black box function, which receives 6 input and outputs Lo (Tp). other words, the task performance of 0' depends on how gives attention when learning evidence text data. And & can be optimized to minimize the (TD). For this, the procedure of (a) and (b) developed to corresponding steps of Eq. and (5) where a and are the learning rates for each respective update.
0' - aV VVTTPPLL (D, WD.))
(4)
o 6 BV v (Tp)
(5)
Algorithm 1: Optimization of Ttain-Attention
W
De,
optimally capabie
:a model
of the task 7a
Require: Dataset 2 {()TT))
6
Require:
Learning rate for 0, Φ ,3
Owz
Co (Tp)
model traned with
Require:
Max iteration step for training 0: M Initialize: base model (f), Train-Attenton (Ø)
f
while é not converged do Sample a data pair (D, (D) ~0 WD.s + ((D)
Figure 3: Optimal W leads f closer to d".
Predict weights
Update ø
for M times do
W
Lieern TTPPLL(DDW.)
Evaluate Vallearn
Data
PPLe
Update: - - a aoo Litaaa
end for
Uodist 8
(0 updated to 0')
Loolue Cgr (TD)
0
E(Task)
Infor
Evaluate VoLaoiue
Train Allention
Update: + Φ & Looouee
Basa Model
Task
Gradient tracking Or adlere trackingg of 0
Reset 0 to initial point 0 end while
Figure 4: One step update of Φ.

---

First, base model 0 is updated to 0' through TWL, with the token weight WD which is generated from ø. Second, the meta-learner Φ is updated based on the task performance, Eg (Tp). These two steps of update can be also interpreted like Figure33 As the model A steps out to a new state 0', the resulting position depends on which token weight (W) ) applied. Some positions are closer to the 0", a model optimally capable of the task TD. as the distance is measured with Lør (TD). We can conclude the weight with a shorter distance (W:)) is more optimal, in the perspective usefulness. To prevent the 0 from converging to the point 0*, which disables the measurement of distances, we reset the model parameters to the initial state 0 after every update of Φ. More detail is depicted in Figure 4 and Algorithm|| [ As the gradients of parameters of & tracked during the updating of 8, its actual implementation is akin to the second derivative. The max iteration step of (denote as M in Algorithm D is fixed to 1 through our experiment. We employ gradient accumulation when updating Φ for batch effect.
On the inference phase, 0 learns data in TWL. manner as in Eq. with the parameter of é frozen. Attention-Augmennn
This system is structure of is fixed to the small model TTnnLLlama-I.IB (Zhang 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 101M 10aM-pareeetrr bidirectional transformer (BERT) Devlin 2018) as a Train-Attention (TA) to further reduce resource requirements.
Language Model (TAALMM). In this work, the foundational
4 Experiment
We conduct experiments on two benchmarks. One is our newly designed LAMA-CKL, and the other the established benchmark. TEMPORALWIKE (Jang et al.) [022]. 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-CKI. benchmark. For T'EMPORALWISS most configurations are aligned with the original work.
4.1 LAMA-CKL.
Train Phase: : To-Learn dataset
Test Phase: To-Learn task Not-To-Forget task
Taskc
Above entranca Barscobo Castle is an armorial panel bearing the arms Macielian and Gardo v with the initials of Willism Maclellan, who buit the castle, and his wite Mary Gordon, the natura daughter Robert Gordon of Lochinvar, 4th
To-Leam task
Lochinvar is <objects
Areway
500
Task: prodict <objacts
no evidonce
Fot-ge-Forgettaa task
language Wales
cobjects
Weish
Franch
Ameayy
30 epochs
Figure 5: Evaluation procedure of the LAMA-CKI benchmark.
For LAMA-CKL, tailor the LAMA dataset (LAnguage Model Analysis (Petromi et al. 2019] to assess the CKL performance, especially the T-REx [Elsahar et al., [018]] part which consists of data from Wikipedia and Wikidata. LAMA is 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 TO-RGET sets from LAMA assess each dimension. During the evaluation, as illustrated on Figure 5) the model learns the evidence documents in the TO-LEARN set. is then tested on both the TO-LEARN task and the NOT-TO-PORGET task to assess plasticity and stability, respectively.
Dataset Setup Here, we outline a protocol for sampling test corpora for the LAMA- -CL bench- mark. 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 variant
5

---

typically requires updates. (2) to ensure the concept of "knowledge new to the model", we select units where the task accuracy zero when measured with pre-update baselines. Conversely, because NOT-TO-FORGES 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 We recommend sampling new dataset by the specified constraints when evaluating models outside of the ILLMA-familly for more accurate assessment. The categorization of fime-variant and time-invariant follows [aang al. [2021]. Each selected unit includes (1) an evidence document, (2) knowledge base triple (e.g.. <Lochinvar, is an instance of, castle>), and (3) a descriptive sentence encapsulating the triple (e.g.. "Lochinvar is 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.1
Evaluation Setup During the evaluation, each epoch consists of both a training phase and a test phase. In the training phase, the model is trained TO-LEARN set. Subsequently, in the test phase, the model's prediction accuracy for the object labels 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 O signifies the model's plasticity. Conversely, 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.
of 500 evidence documents from the
In the proposed benchmark LAMA-CKL, we suggest four main factors as evaluation indicators. 1) Top Ace: the highest TO-LEARN accuracy among checkpoints of 30 epoch. 2) Epoch: the epoch where the Top Acc appears. 3) NF Acc: NOT-TO-FORGE accuracy of the checkpoint model which is the same Top Ace. 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 better if higher, while factor 2 is better if lower. The detailed configurations for training datasets are in Appendix A.1.1
Train-Attention Training Setup Referring to Algorithm T the training procedure of Train- Attention requires data 'D and related task Tp. For LAMÁ-CKL dataset, we assign evidence document of each unit as D. and knowledge base triple in a document of schematic form as TD. The perplexity of object token is assigned as the objective of Φ. Figure 6jshows the heat map of ooken-weight that
generates. Train-Attention seems to give more attention to entities of certain categories, rather than words with general grammatical roles. We describe the detailed configuration and findings on the training of Train-Attention in Appendix
A.2 A.2
The
perfoct
gane
obo
race
Bun
thrd
peat
a'to
Young
and
Add
Etazo
perlect
game
und
addicral ED
Sand
Jahneon
Mark
Hal
pitsh
WTE
pich
ed
hi
additoonll no
1307
Fhe
Nattonal League
Figure 6: Heat map of token weights from Train-Attention. Orange color indicates higher weights.
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. [022]] standard finetune, K-Adapter [Wang et al. 2020). Mix-review (He et al. 2019].
LoRA [Hu et al. 2021]. RecAdam [Chen al., 2020]. We regard standard finetune on QLoRA as 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 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 srper-parameter following the optimal of the original. We also evaluate a model trained in TWL manner with Oracle token weight. For which, weight of exclusively assigned to the object label token in the evidence document, and the rest is assigned zero weight. Oracle is compared for two purposes: (1) To prove the concept that token-weighted learning has the advantage for CKI.. (2) To check the performance upper bound of Train- Attention. Detailed configurations are in Appendix A.3

---

To-Leam Accuracy (7B)
Not-in-Forget (78)
0.4
0.90
Fintuneiggora)
accuracy
KAdaptar
Mix-reviewtolorr)
0.85
Recddamiqoorl
1/q/cral
TAALMIclora)*
#netunelofore)
KAdapter
Muxreviewiclora)
Rockdamigloral
0.70
RHD-Ilalora
TAALM(C(ara)
10
15 epoch
20
25
30
10
15
25
30
epoch
Figure 7: LAMA-CKL performance large (Ilama2-7B) baseline models. The graph on the left represents TO-LEARN task, and the graph the right represents NOT-TO-FORGET task performance. The x-axis is the learning epoch, and the y-axis is accuracy.
Table LAMA-CKL performance of Llama2-7B based baselines. The evaluation indicator of each column is explained on $4.1 The best performance is marked as bold while the second best is underlined.
Top Acc
Epoch
NF Acc
Total Knowledge
Fioetune(QLoRA)
0.1150
16
0,8174
0.9324
K-Adapter
0.1260
30
0.9980
1.1240
(QLo)A))
0.0800
25
0.7988
0.8788
RecAdam(QLoRe)
0.1000
24
0.7933
0.8933
RHO-I(QLORA)
0.1410
18
0.8223
0.9633
TAALM(QLORA)
0.4290
4
0.8983
1.3273
4.1.1 Result & Analysis
TAALM achieves substantial CKL performance As results show in Figure [ and Table [ our method (TAALM) overwhelms other baselines ability and speed of learning. Ours records the Top Ace 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, that it is meaningless to compare K- 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 bigger capacity for learning and less forgetting than standard finetuning, However, the advantage is minimal compared to ours.
conclude
Combination with TAALM improves all of the previous baselines As Train-Attention is approach to manipulating loss values on the end side, it is casily 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 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 assumed that TWL with Oracle label is the upper bound of performance of TWL with
because it is sharply targeting only
7

---

Lo-m
Not Forget Accuracy (78)
1.5-
5.9-
0.
Freturalckral WALMMaOa""
5.8
0.1
westendtebball Grade
5.75
15
epoch
spech
Figure 8: Comparison between Oracle, standard finetuning, and ours, tested on LAMA-CKL.
the necessary token. Referring to Figure 17, for Oracle, Top Ace 0.5470, epoch 17, and NF Acc is 0.9002. Oracle shows that Top Acc is 4.75 times higher and NF Acc is 1.11 times higher than standard finetune, thus proving the substantial advantage of weighte learning on CKL. Also, TAALM nearly approaches Oracle, as it achieves 78.2% of Oracle Top Acc. NF Acc of ours also maintains a similar level to Oracle, indicating that Train-Attention is optimized close to the upper bound. It also indicates that optimization through meta-learning could excel human labeled weight.
Small (1B) TAALM excels large baselines We also experiment the smaller (TinyL.ama-- 1B) baselines, and TAALM on IB records the best compared to 7B baselines. This observation indicates that our method outperforms other baseline methods even with significantly smaller parameter sizes and computational resources. Related Table and Figure are on Appendix A.4.1
TA on BERT achieves comparable performance with small resources We train BERT as a TA and evaluated 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 (I.1B) TA (Appendix D)
Ablation study on various design choices 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
4.2 TEMPORALWIKI
We experiment on the original CKL learning benchmark TEMPORALWIKI [Jang et al. 2022). where models have continually learn Wikipedia documents of serial periods (0809, 09T0, 1011, 1112) and test on the corresponding TWWIKI-PRBES. which is dataset of knowledge base triples. As we train Train-Attention on the 0809 data, tests are conducted on the rest. We experiment with only small (TinyLlama- IB) 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-CKI on 1 We only consider training DIFFSET, which is the only changed part of Wikipedia, because 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
4.2.1 Result & Analysis
Referring to Table 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 aligns with the L-CLL 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 the TEMPORALW IKL, referring
8

---

Table 2: TEMPORALWIK. performacne of small (TinyLLama-1B)) baselines. Un refers UNCHANGED,
C refers CHANGED, Avg refers the average of the two. TAALM is our method.
QLoRA based
TWiki-Probes-De
Siiki-Probbs-1ue 1011
TWiki-Probes- 1112
Un
Avg
C
Avg
Un
Avg
Qoo)tuncQQLoRA)
9.999
10.057
10.028
9.554
9.531
9.543
9.736
9.632
9.684
(ixxrrevic((QooA))
9.529
9.579
9.554
9.514
9,486
9.501
9.562
9.452
9.507
RecAdam(QLoRA)
9.514
9.604
9.559
8.992
9.031
9.012
9.579
9.479
9.529
RHO-I(QLORA)
4.389
4.624
4.507
4.360
4.395
4.3775
4.471
4.717
4.594
TAARM
4.019
4.268
4.1435
4.030
4.154
4.092
4.036
4.357
4.197
(b) K-Adapter based
TWiki- Probes-091 0910
WWiki-Probes-ie 1011
TWiki- Siiki-Probss-t 1112
Un
C
Avg
Un
Avg
Un
C
Avg
Finetune(K- Baattune(K-aaappp Adapter)
2.768
2.982
2.875
2.598
2.679
2.639
2.552
2.669
2.611
Mix-review Adapttr))
2,486
2.724
2.605
2.334
2.446
2.390
2.284
2.425
2.355
RecAdam(K-Aduppe Addatter)
2.494
2.710
2.602
2.323
2.415
2.369
2.248
2.375
2.312
RHO-1(K (Ada-t())
2.722
2.950
2,836
2.538
2.634
2.586
2.478
2.603
2.541
TAALM (K-Adapter) on LAMA-CKL
2.364
2.601
2.483
2.168
2.282
2.225
2.123
2.307
2.215
TAALM (K-Adapter)
1,980
2.194
1,9705
1.907
2,034
2.087
1,901
2.070
1.986
Table 2b
It achieves the second-best performance, indicating robustness across different distributions
of tasks
4.3 Why LAMA-CKL: clear contrast of plasticity and stability
For the benchmark TEMPORALWIKI because DIFFSET corpora of evidence documents for CHANGED set, learning of DIFFSETS is supposed to result in performance improvement CHANGED set and forgetting of UNCHANGED set. However, during our experiments, we observe that both CHANGED and UNCHANGED performance tend 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 DIFESET contains evidence documents for both the CHANGED and the UNCHANGED sets (Appendix C.4). This is complicating factor the evaluation of stability. Second, the experimental setup involves training on vast amount of data, an average of 707K documents per period, for just 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 closer to continual learning setup rather than CKL. To address this issue, we structured the LAMA-CK. A-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- pre-update baselines. (2) We conduct training over multiple epochs on relatively small dataset observe acquirement of knowledge. In practice, our benchmark shows clear upward trend in the TO-I.FARN set and a distinet decline the NOT-TO-FORGET set as training progresses, clearly demonstrating the contrast between plasticity and stability.
5 Conclusion and Limitation
In this paper, we demonstrate that the application of Train-Attention significantly enhances CKI. 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 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 data-task pair. there no paired dataset, it can be doubted that training is difficult. However, every knowledge has its purpose, and we can find strategy to discover it. 1) Search: When the data and task pools are separate, we can join highly related pairs via searching. TEMPORALWIH
also a dataset which data and tasks are not paired, thus 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. Synthetie tasks can be generated via instruction-tuned L.M,, 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 increase the ability of CKL. therefore enhancing the prac- ticability of LLMs and saving the computational resources for fine-tuning entire huge I.L.Ms. We believe that this paper does not have any immediate negative societal impact.
Acknowledgement
This work supported by STEAM R&D Project, NRF, Korea (RS-2024-00454458) and Institute of Information Communications Technology Planning & Evaluation (IITP) grant funded by the Korean goverament (.STT)NNo.R--200--1D University)) and (2022-0-00077, Extraction, Reasoning, and Inference from Heterogeneous Data). Jinyoung Yeo and Dongha Lee the co-corresponding authors.
361, Artificial Intelligence Graduate School Program (Yonsei Technooogy Development for Commonsense
References
R. Aljundi, F. Babiloni, M. Elhoseiny, M. Rohrbach, and T. Tuytelaars. Memory aware synapses: Learning what (not) forget. In Proceedings of the European conference computer vision (ECCV), pages 139-154, 2018.
M. Andrychowicz. M. Denil, S. Gomez, M. W. Hoffman, D. Pfau, T. Schaul, B. Shillingford, and N. De Freitas. Learning to learn by gradient descent by gradient descent. Advances neural information processing systems, 29, 2016.
T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-19011 2020.
Chen, Y. Hou, Y. Cui, W. Che, T. Liu, and X. Yu. Recall and learn: Fine-tuning deep pretrained language models with less forgetting. arXiv preprint arXiv:2004.1265) 12651,, 2020.
T. Dettmers, A. Pagnoni, A. Holtzman, and L. Zettlemoyer. Qlora: Efficient finetuning of quantized 1lms. Advances in Neural Information Processing Systems, 36, 2024.
J. Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:/810.0480)) 2018
H. Elsahar, P. Vougiouklis, A. Remaci, C. Gravier, J. Hare, F. Laforest, and E. Simperl. T-rex: A large scale alignment of natural language with knowledge base triples. In Proceedings of the Eleventh International Conference Language Resources and Evaluation (LREC 2018), 2018
C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pages 1126-1135. PMLR, 2017.
L. Franceschi, P. Frasconi, S. Salzo, R. Grazzi, and M. Pontil. Bilevel programming for hyperpa- rameter optimization and meta-learning. In International conference on machine learning, pages 1578-1577. PMLR, 2018
10

---

T. He, J. Liu, K. Cho, M. Ott, B. Liu, J. Glass, and F. Peng. Mix-review: Alleviate forgetting in the cectrain-finttur framework for neural language generation models. 2019.
T. Hospedales, A. Antoniou, P. Micaelli, and A. Storkey. Meta-learning in neural networks: A survey. IEEE Iransactions on patlern analysis and machine intelligence, 44(9):5149-516) 2021.
L. Hou, R. Y. Pang. T. Zhou, Y. Wu, X. Song, X. Song, and D. Zhou. Token dropping for efficient bert pretraining. arXiv preprint arXiv:2203.132400 2022.
N. Houlsby, A. Giurgiu, S. Jastrzebski, B. Morrone, Q. De Laroussilhe, A. Gesmundo, M. Attariyan, and S. Gelly, Parameter-efficient transfer learning for nlp. In International conference machine learning, pages 2790-27222 2799. PMLR, 2019.
E. J. Hu, Y. Shen, P. Wallis, Z. Allen ZZu, Li, S. Wang. L. Wang, and W. Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2 09685, 2021.
J. Jang, S. Ye, S. Yang, J. Shin, J. Han, G. Kim, S. J. Choi, and M. Sco. Towards continual knowledge learning language models. arXiv preprint 0rXiv:21100032111 2021.
J. Jang, S. Ye, C. Lee, S. Yang, J. Shin, J. Han, G. Kim, and M. Seo. Temporalwiki: A lifelong bench- mark for training and evaluating ever-evolving language models. arXiv preprint arXiv:2204.1221 2022.
R. Kemker, M. McClure, A. Abitino, T. Hayes, and C. Kanan. Measuring catastrophic forgetting in neural networks. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
J. Kirkpatrick, R. Pascanu, N. Rabinowitz, J. Veness, G. Desjardins, A. A. Rusu, K. Milan, J. Quan, T. Ramalho, A. Trabska-Barwinnka et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, (14(13)35211335 2017.
M. F. Land and S. Furneaux. The knowledge base of the oculomotor system. Philosophical Transactions of the Royal Society of London. Series B: Biological Sciences, 332(1359); 1997.
Z. Li, Zhou, F. Chen, and Li. Meta-sgd: Learning learn quickly for few-shot learning. arxiv 2017. arXiv preprint 1rXiv:1707.0900 2017.
T.-Y. Lin, Goyal, R. Girshick, K. He, and P. Dollár. Focal loss for dense object detection. In Proceedings of the IEEE international conference on computer vision, pages 2980-2988, 2017.
Z. Lin, Z. Gou, Y. Gong, X. Liu, Y. Shen, R. Xu, C. Lin, Y. Yang, Jiao, N. Duan, al. Rho-1: Not all tokens are what you need. arXiv preprint arXiv:2404.07965 2024.
D. Lopez-Paz and M. Ranzato. Gradient episodic memory for continual learning, Advances in neural information processing systems, 30, 2017.
F. Petroni, T. Rocktäschel, P. Lewis, A. Bakhtin, Y. Wu, A. H. Miller, and S. Riedel. Language models as knowledge bases? arXiv preprint (ttl::/9990000000 2019.
C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-67 2020.
D. Rolnick, A. Ahuja, J. Schwarz, T. Lillicrap, and G. Wayne. Experience replay for continual learning. Advances in neural information processing systems, 32, 2019.
V. Sanh, A. Webson, C. Raffel, S. H. Bach, L. Sutawika, Z. Alyafeai, A. Chaffin, A. Stiegler, T. L. Scao, A. Raja, et al. Multitask prompted training enables zero-shot task generalization. arXiv preprint arXiv: 10.08207 2021.
H. Shin, K. Lee, J. Kim, and Kim. Continual learning with deep generative replay. Advances in neural information processing systems, 30, 2017.

---

F.K. Sun, C.-H. Ho, and H.-Y. Lee. Lamol: Language modeling for lifelong language learning. arXiv preprint arXiv:199.0333 2019.
H. Touvron, Martin, K. Stone, Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint Xiix:3307009288 2023.
R. Wang, D. Tang, N. Duan, Z. Wei, X. Huang, G. Cao, D. Jiang, M. Zhou, et al. K-adapter: Infusing knowledge into pre-trained models with adapters. arXiv preprint arXiv:2002.01808, 2020.
Wei, M. Bosma, V. Y. Zhao, Guu, A. W. Yu, B. Lester, N. Du, A. M. Dai, and Q. V. Le. Finetuned language models are zero-shot learners. arXiv preprint arXiv:2109.01653 2021.
F. Zenke, B. Poole, and S. Ganguli. Continual learning through synaptic intelligence. In International conference on machine learning, pages 3987-3995. PML.,, 2017.
Zhang, G. Zeng, T. Wang, and Lu. Tinyllama: An open-source small language model. arXiv preprint ar Xix::2401.002883 2024.
12

---

A LAMA-CKL Benchmark Additional Detail
A.1 Dataset setup detail
As each unit of the LAMA dataset contains multiple evidence documents, we sample one document for each unit, based on the following criteria: (I) The document should exceed length of 70 tokens. (2) The document must include both the subject and object entities. These documents are then truncated maximum length of 512 tokens. We assess the accuracy of the knowledge base triple of each unit to filter them. We utilize the TinyLlama-I. 1.1 Zhang et al. [2024] and Llama2-7B Touvron
2023) models, integrated with QLoRA Dettmerss al. [2024] and K-Adapter Wang et al. [020]. total 4 variations. Units with unanimous accuracy among the four models are only selected TO-LEARN or NOT-TO-FORGET set. Because NOT-TO-FORGET accuracy differs among descriptive and schematic form queries, NOT-TO-FORGET set for descriptive and schematic tasks are separately sampled, allowing overlap. The distribution relation categorics among TO-LEARN and NOT-TO-FORGET is depicted in the Figure 9 Also the relation categories of each TO-LEARN and NOT-TO-FORGET are detailed in Table
Table presents data statistics for NO1-TO-FORGET and
TO-LEARN sets.
Not Forgetidesc/pttee)
P13
123
100
127
100
P30
1900
P100
R127
15:30
0150
PLTS
1178
1776
1304
1407
(413
ruGs
ms50
1740
1000
JodhB.forgelischemaaaa
300
F17
F33
122
122
1500
1903
1500
P109
P227
P131
P130
11.11
1940
1350
15.5
1179
1200
1001
1275
1279
1901
1964
1907
(422
1963
1507
res0
1940
ms00k
11378
To-Learn
190
122
1206
P108
m227
m140
P150
PL78
P120
P264
P364
IN0
[413
M527
mo30
10003
PLaz
Train-Attention train dataset
200
REU
mn
P/00
RZ4
07/00
2003
93376
2002
Figure 9: Distribution of relation categories of each dataset in LAMA-CKL
Table 3: LAMA-CKA data statistics.
Size
Avg. Evidence token
Avg. Label token
NOT-TO-FORGET (descriptive)
500
141.1
1.5
-ORRGET (schematic)
500
153.4
1.9
TO-LEARN
500
112.8
1.0
Train-Attention Train Data
4166
136.0
1.1
'Directly referred from 'Relations of InvariantLAMA of Jang al. 2021
13

---

Table 4: Relations of each NOT-TO-FORGER and TO-LEARN set.
(a) Relations of NOT-TO-FORGE
Relation Code
Template ([X), [Y])
Relation
P19
[X] was born in [Y]
place of birth place of death subclass of official language original network shares border with named after
P20
[X] died in [Y].
P279
[X] is subclass of [Y].
P37
The official language of [X] [Y].
P449
[X] was originally aired [Y].
P47
[X] shares border with [Y].
P138
[X] is named after [Y].
P364
The original language of [X] is (Y).
original language of film or TV show
P527
[X] consists of [Y].
has part manufacturer
P176
[X] is produced by [Y].
P27
[X] [Y] citizen.
country of citizenship
P407
[X] was written in [Y].
language of work or name continent developer capital of located in the administrative territorial entity
P30
[X] located in [Y].
P178
[X] is developed by [Y].
P1376
[X] is the capital of [Y],
P131
[X] is located in [Y].
P1412
[X] used communicate in [Y].
languages spoken, written or signed country location
P17
[X] is located in [Y].
P276
[X] located in [Y].
P937
[X] used to work in [Y].
work location
P140
[X] is affiliated with the [Y] religion.
religion native language twinned administrative body applies to jurisdiction country of origin capital location of formation part of
P103
The native language of [X] is [Y].
P190
[X] and [Y] are twin cities.
P1001
[X] is legal term in [Y].
P495
[X] was created in [Y].
P36
The capital of [X] is [Y].
P740
[X] was founded in [Y].
P361
[X] is part of [Y].
(b) Relations of TO-LEARN
Relation Code
Template ([X), [Y)
Relation
P31
[X] is a [Y]-
instance of position held field of work occupation employer owned by
P39
[X] has the position of [Y].
P101
[X] works in the field of [Y].
P106
[X] is [Y] by profession.
P108
[X] works for [Y]
P127
[X] is owned by [Y].
P136
[X] plays [Y] music
genre headquarters location
P159
The headquarter of [X] is in [Y].
P264
[X] is represented by music label [Y].
record label
P413
[X] plays in [Y] position. [X] is member of [Y].
position played team speciality member of diplomatic relation instrument
P463
P530
[X] maintains diplomatic relations with [Y].
P1303
[X] plays [Y].
A.1.1 Evaluation setup detail
Hardware and hyper-parameter: During training TO-LEARN documents, 8 RTX 3090 GPU (24GB) are used, with global batch size of 64. A total of 30 epochs took 25 minutes of GPU time. Learning rate lc-4, AdamW optimizer, and max length of 512 tokens are applied.
14

---

Accuracy measurement Accuracy is utilized as a metric (Top Acc, NF Acc) to assess the efficacy of models in the task of label prediction. This metric quantifies the proportion of label tokens correctly identified by the model out of the total label tokens presented.
A.2 Train-Attention training detail
Dataset We select LAMA units from both time-variant and time-invariant sets, specifically those with an accuracy below 0.5, while ensuring no overlap with the NOT-TO-FORGET and TO-LEARN sets. Furthermore, the initial data exhibit a disproportionately large number of units categorized under relation P530 ('diplomatic relations of the country"). ). address this imbalance, we adjust the frequency of P530 units to match that of the second most prevalent type. Consequently, the final version of the Train-Attentic training dataset comprises a total of 4166 units, detailed in Table 3
We propose the schematic form to arrange the knowledge base triple into a human- readable sentence. fits the one-directional feature of causal LM, as opposed to the descriptive form where important information sometimes appears behind the label tokens. The template and example are on Appendix [ We employ schematic form in training Train-Attention, while employing descriptive form during evaluation.
Training We utilized small (TinyL and large (Llama2-7B) models from the LLaMA family, integrated with QLoRA [Detumers al., 2024]. We found that using large rather than small model for the base model (0) results in faster convergence. On the other hand, employing a large or small model for the Train-Attention model (ø) made little difference in the aspect of the convergence step and validation score. So we adopt large model for & and a small model for é while training In the test phase, We find that 0 and Φ are still compatible even though they don't t share the same background model, as long as they use the same tokenizer. Therefore we utilize small-size Train-AAtention for both small and large baseline experiments.
We initialize the parameters of the Train-Attention head, which is a decoder layer as depicted in Figure 2h, using a normal distribution with mean of 0 and standard deviation of 0.0001. Consequently, the weights generated by the initialized Train-Attention are equivalent to the uniform weight where all w; = This approach adopted observe the convergence of the Co (Tp) more clearly, ensuring that the loss value declines from the initial state if training proceeds normally.
A single A100 (82GB) GPU is used, and the effect of batch size 16 is achieved through gradient accumulation. A checkpoint of 400 global steps is used (takes about 6 GPU hours). Of the 4166 train data, 100 are used for validation. The learning rate of 2e-4 and Adam W optimizer is employed for both the base model and Train-Attenticn
A.3 Baselines detail
In this section, we describe the detailed configuration of QLoRA and K-Adapter. The other baselines follow the previous CKL work of (Jang et al. 2022
OLoRA Hyper-parameter, of QLoRA follow one of the optimal of the original work (Dettmers)
ct al. 2024]. We employ LoRA r = 64, a = 16, NF4 with BF16 computation datatype. A total of 160M parameters expanded for the large (Llama2-7B) baselines.
K-Adapter papper-prramters of K-Adapter follow the previous CKL work of Jang et al. 2022 total of 303M parameters are expanded. The trainable parameters are double the QLoRA. Efforts are made to reduce the learnable parameters of K-Adapter, but we observe that the current settings are the minimum level necessary for K-Adapter to function effectively. In experiments, K-Adapter demonstrates greater stability than QLoRA, which could be partially attributed to its larger parameter size. We also employ the same quantization configuration of QLoRA for K-Adapter, NF4 with BF16 computation datatype.
15

---

A.4 Experimental results additional detail
A.4.1 LAMA-CKL Experiment on small (1B) baselines
Figure 10land Table 5] describe performances small (TinyLlama- baselines on the LAMA-CKL. The performance of TAALM based on the small (TinyLlama- 1B) architecture is significantly better even when compared the large (Llama2-7B baselines.
To-Leam Accuracy (18)
Not-To-Forget Accuracy (18)
1.00
0.95
0.25
0.20
Finetuneqporrll
KAcapter
0.85
RecAdamielora)
RHD lielar)
0.00
TAALMigiora"
Inntuneigloral
Niw.eviewiggorrl
0.05
RecAdamiclora)
0.70
RHD (icoor))oor))
0.00
TAALMalora!"
20
20
epoch
egoch
Figure 10: LAMA-CKI, performance of small (lany-.lama-1)) baselines.
Table LAMA-CKL performance of small (TinyLlama- baselines.
Top Acc
Epoch
NF Ace
Total Knowledge
Finetunc(QLoRA)
0.0700
29
0.7693
0.8393
K-Adapter
0.0240
22
0.9908
1.0448
Mix-review(QLoRA)
0.0600
23
0.8161
0.8761
BecAdam(QLoRM))
0.0600
29
0.7719
0.8319
(QLOR())ORA)
0.0900
13
0.8523
0.9423
TAALM(QLORA)
0.3260
4
0.9078
1.2338
A.4.2 Combination of ours (TAALM) and other baselines
Figure I presents that the performances of all baselines significantly improved when combined with the Train-Attention, indicating synergistic compatibility of our method.
Referring to Table 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-Attentic not only to the train data but also to the review data. As result, it 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-Attentio 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, 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 minimal trade-off between updating and forgetting.
16

---

Mix-review(qlora)
RecAdamlqlora
Adapter
10-
30-
[.61
20
1.0 10
10 25
15 20
25
Figure 11: Comparison of each baseline alone and combined with our method. Each title on the plot represents the baseline method. The gray line represents the baseline alone, and the red line represents the combination with TAALM. Solid line for TO-LEARN, dashed line for NOT-TO-FORGET. All are based on Llama2-7B, and tested on LAMA-CKL.
Table 6: Combination of ours (TAALM) and other baselines. Based Llama2-7B, tested LAMA-CKL.
Top Acc
Epoch
NF Acc
Total Knowledge
ours K-Adapter
0.3320
21
0.9747
1.3067
ours(QLoRA)) + Mix-review
0.4500
3
0.9012
1.3512
ours( QLoRA) + RecAdam
0,4360
7
0.8982
1.3342
ours K-Adapter RecAdam
0.2640
13
0.9730
1.2370
ours K-Adapter + Mix-review + RecAdam
0.3140
25
0.9677
1.2817
ours QLoRA)
0.4290
4
0.8983
1.3273
Te-Leam Accuracy (78)
Not-lo-Forget Accuracy (78)
1.000
0.4
0.975
0.350
0.3
0.2
K-Adapter
KAAdaper
oursigloral Mis-reviww
0.1
ours/gloral Recidam
ourstglora) RecAdsm
K-Acaptor
nurs. KAddapeer RecAdam
ours K-Aapter Mix-roviow RocAdam
KoAapt Masstecee RecAdam
0.0
oursiglorel
10
10
Figure 12: Combinations of other methods with ours (TAALM) Llama2-7B base model, tested on LAMA-CKL.
A.4.3 Rationale for the accuracy drop
In the main experiment (Figure 7). the accuracy declines after rapid peak after about the fourth epoch. We hypothesize this is result of overfitting. While the token weights from TA include beneficial targets, there must also be some that 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 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 's performance does not decline with further training, the differences in accuracy trends could be caused by false targets. This phenomenon can also 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 common due overfitting, this cycle seems to occur more rapidly for TAALM.
17

---

1.0
0.9
0.9
Forget
MP
Finetunelciora)
0.7-
<Acaaber
0.7
Reckdamiqlora)
Mix reviewewwl)))))
Reckdamiglor))
1iciora)
RHD Jielors
ours KAAdapter
Mix-review
cursiolorat Mivieview
oursiqlora) RecAdom
cursiglora) RecAdam
KAdapter RecAdam
curs KiAdaater Rechdam
K-Aapter Mix-review
curs KAapaar Nixrevias RecAcom
oursiqiora)
oursigloret
0.5
0.0
0.0
0.1
0.2
0.3
0.4
0.5
Accuracy
(a) Dots
(b) Lines
Figure 13: The trade-off between plasticity (TO-LEARN) and stability FOORGE) visu- alized, for all baselines including combinations with our method. All models are based on large (Llama2-7B) architecture. Gray dashed lines stand for the zero-sum state where the Total Knowledge (sum of the performances of TO-LFAAN and NOT-TO- NOT-TO-FORGET) 1. The right and upper sides of the gray lines indicate the more efficient system where learning causes less forgetting. (a) presents only the checkpoints of the Top Acc, while Lh)L presents whole checkpoints of 30 epochs as lines.
A.4.4 Reporting error ranges of the main experiment
We conduct five independent runs using different random seeds the training data-loader, and display the results with error ranges at 2o in Figure 14 This demonstrates the consistent performance of our method. All main experiments presented use the data loader with random seed 42.
Accuracy (7B)
Forget Accuracy (7B)
1.0
0.3
Fhotunelglera)
sccuracy
KAdapter
accuracy
Mix-revieofolora)
Finetunelglore)
0.2
RecAcamigiora)
Kddapter
liqicral
0.4
Mix-neviewtolora)
TAALMiolara)*
RecAdamiqioral
RHD 1igloral
TAALMa
0.0
0.0
10
15
20
30
10
20
30
epoch
opoch
(a) Dots
(b) Lines
Figure 14: Performance of the large-scale baseline models (Llama2-7B 7B) on LAMA-CKL. depicted with #20 error ranges. Results calculated from five random trials, employing different random seeds over the train data-loader.
18

---

B Descriptive Form & Schematic Form
Each unit of the LAMA dataset comprises knowledge base triple and corresponding sentence that encapsulates the information contained within the triple. These sentences are presented in short dexcriptive 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 mnidirectional nature of causal language models (LMs). In the descriptive form, critical information often comes after the label tokens, which can be problematic. For instance, 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 the train-attention learning process.
Knowledge Base Triple
<Subject>, <Relation>, <Object>
Descriptive Form
Schematic Form
Template
Template
Example
The <Relation> of <Subject> is <Object>
Guess the Object.
Guess the Object.
Subject <Subject>
Subject Wales
Example
Relation <Relation>
Relation official language
The official language of Wales is Welsh
Object <Object>
Object Welsh
Descriptive form and our proposed schematic form template rearrange knowledge triple Figure 15:
into a human readable document.
C Detail of TEMPORALLWKKI Experiment
C.1 Corpus setup
The original corpus consists of two datasets to learn for 4 periods (0809, 0910, 1011, 1112); (1)
Wikipedia: the whole snapshot of Wikipedia documents of each period. (2) DIFFSETS: the segments of documents that only contain changed information compared to the last period snapshot. The corpus also includes test datasets (TWiki-Probes) corresponding to each period. Just like the LAMA-CKL, the Wikipedia snapshots and DIFFSETs consist of evidence documents, while TWiki-Prodee Probes consist of knowledge base triples. Because we used the data from the very first period train Train-Attention only the data from the next three periods (0910, 1011, 1112) are used for the evaluation.
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 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 supposed to measure plasticity, while UNCHANGED set is supposed to measure stability.
C.2 Evaluation and baseline setup
We follow evaluation settings of the original work et al. [202]] a model continually learns train datasets of each period while evaluating performance on TWiki- after training of each period. The evaluation metric the mean perplexity object label words. One dataset is updated for only 1 epoch.
19

---

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 LAMA-CKI,
Baselines and hyper-parameter settings follow that of LAMA-CKL al. 2022]
$4. which follow the Jang
C.3 Train-Attention training setup
The DIFFSET and Twiki-Pres of the first period (0809) used for training Train-Attention. To train
(Algorithm T. data (D) and the corresponding task (Tp) must be paired, while the documents in DIFFSET and knowledge base triples in Twiki-Prohes 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
not paired. To tackle
C.4 Analysis the results
In this section, we present the detailed experimental result on TEMPORALWIK,, specifically involving the discussion in $4.3 the reason propose the benchmark LAMA-CKL. Figure 16 indicates that performances in both Changed and Unchanged mostly move in the same direction when the baselines trained the DIFFSET of each period, presenting plasticity (Changed) and stability (Unchanged). This is contrasting to the observation on our new benchmark, LAMA--KL (Figure 13), which exhibits clear trade-off between plasticity (TO-LEARN) and stability (NOT-TO-FORGET). This suggests that the LAMA-CKL is better benchmark for observing the trade-off between plasticity and stability.
trade-off between
Table presents the portion of task sentences 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 DIFFSET partially results in improvement in the performance the UNCHANGED set. 12) 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.
The Unchanged set from the last period (TWiki-Probes
Table 7:
The portion of units in each period of TWiki-Probes, which have corresponding evidence document in the DIFFFSET. Un and indicate Unchanged and Changed set, respectively.
TWiki-Probes-0910
TWiki-Probes-Ile 1011
TWiki-Probes- 12
Un
C
Un
Un
C
DIFFSET-0910
7.8%
83.6%
7.6%
39.5%
7.7%
25.3%
DIFFSET-1011
7.4%
81.8%
7.9%
26.2%
DIFFSET-1112
8.7%
83.3%
20

---

QLoRA based
K-Adapter based
10
Finetune
4.5
Finetune
Mix-review
RecAdam
RecAdam
TAALM
TAALM
4.0
of Unchanged
PPL Unchanged
3.5
3.0
2.5
5-
2.0
5
10
2.0
2.5
3.0
PPL of Changed
PPL Changed
Figure 16: The experimental result of smyll(TinLllmaa13 baselines on TEMPORALWIKI bench- mark. The x-axis and y-axis correspond to the perplexity of Changed and Unchanged sets, respectively. Each color represents baselines, and the arrows represent the experimental result from each period. The start point of the arrow indicates the performance before training, while the endpoint indicates performance after training.
D Reduction of resources through TA on BERT
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) body for TA, which has bigh inferential capabilities even at very small size (108M) compared to the previous body (Tinyllama I.IB), due to its bidirectional property. Since BERT has different tokenizer from our generation model, the Llama family, we integrate BERT with the Llama2 tokenizer and pre-trat pre-train for 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 only single 24GB GPU, significantly reducing resource use compared to the previous version (single 82GB GPU), thus making affordable for the general environment. On the inference, the TA 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.
Table 8: LAMA-CKL performance of small (TinyLlama- 1B) baselines.
(a) Baselines with large generation model (Llama2 7B)
Parameter size of TA
Top Acc
Epoch
NF Acc
Total Knowledge
Finetune
NA
0.50
0.8174
0.9324
TA (Llama)
L.B
0.4290
4
0.8983
1.3273
TA (BERT)
108M
0.3210
0.9388
1,2598
(b) Baselines with small generation model (Tinyllama IB)
Parameter size of TA
Top Ace
Epoch
NF Acc
Total Knowledge
Finetune
NA
0.0700
29
0.7693
0.8393
TA (Llama)
L.IB
0.3260
4
0.9078
1.2338
TA (BERT)
108M
0.2440
9
0.9267
1.1707
21

---

E
Ablation
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.
(I) "roken-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 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 assign lower weight to un un-useful tokens.
Results Known token masking does not yield better results compared to TAALM w/ token- importance weight. We hypothesize that the effect of known masking 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 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.
Table 9: LAMA-CKL performance of Llama2-7B based baselines.
Top Ace
Epoch
NF Ace
Total Knowledge
Finetune
0.1150
16
0.8174
0.9324
TAALM w/ token- iiken-importance weight (ours)
0.4290
4
0.8983
1.3273
TAALM w/ known token masking
0.3920
4
0.9075
1.2995
TAALM w/ token weight dropping < 0.5
0.4100
7
0.9148
1.3248
TAALM w/ token weight dropping 0.8
0.3850
4
0.9267
1.3117
To-Leam Accuracy (7B)
Not-To-Forget Accuracy (7B)
1.00
TAALM w/ docan-importane weight (ours)
TMALM wi known token mesking
0.95
TAALM weight dropping
TAALM w/ toxan waight dropping 0.8
0.90
scouraty
accurscy
0.85
0.80
Fretune
TAALM fikan-importance weight
0.75
TAAL w/ known taken masking
TAALM wi token weight dropping < 0.5
token weight drooping 0.0
10
15
20
25
30
10
15
25
apoch
Figure 17: Ablation study on various design choices with token importance from Train- Atten-nttttionn (TA). TA is Tinyllama- IB, and the generation model is Llama2-7B.
22

---

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. 'he 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 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 (c.g., birthplaces, location of the workplace).
We also provide an attention map of TA trained on the qulti-session chat (Figure 18ch Here, we regard prior dialogue sessions as data (D) and an understanding of the next session as task (Td). 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 's name.
Scot
bore
song
Witer
pulr
Irand
sanger H
player
playet
primutty
and
0
12tt
eto
Anderson
iny
Move
Nich
ttal
(a) Heat map of token weights from Train-Attentee trained on LAMA-CKL.
Origin
fuo
Argondan
Dalay
Dazi
agu
Uragure
20comn
narurt
Chraa
Indomns
Japan
Namm
Phalppiny
Singacoe
New
Guines
atand
Statnam
Zwabrdd
Chrstrua brand
Australa
Pactic
Africa
frund
Cam
41
Foest
Qh
37a
Guree
standd
ervb
l0ue
Nav
Union
ands
Soo
Soch
AYka
and
reb
(b) Heat map of token weights from Train-Atteniin trained on LAMA-CKL.
James
cird
y0u
kmng
been
volumbe
Jemes
been
unturte
exing
yeam
Nice
and
other
JaYeS
What
210
doing
The
The
erl
cione
queen
I0c
ataid
Day
this
Me
and
on
cinc
always
Junes
T16
sounds
Jun
VoU
hov
ary
othor
Mell
country
Ivo
dags
Dal
CY90
What
about
mb
whol
Irg
Mo Ah
Fet
scree
ongineer
guase
(c) Heat map of token weights from Train-Attention trained on Multi-Session Chat dataset.
Figure 18: Heat map of token weights from Prain-Attention. Orange color indicates higher weights.
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The paper includes our mathematical formulation and quantitative expeerimntteeeeeeiieetttt results that reflect and justify the claims in our abstract and introduction.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitation section contains discussion of our method's limitations.
Guidelines:
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: The paper does not include theoretical results.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We show how our meta-learning algorithm be applied to the CKL benchmarks. We clearly explain our algorithm, and architecture with visual aids. We provide anonymized code for our quantitative experiments alongside clear instructions (README.md) for training and evaluation.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide anonymized code for our quantitative experiments alongside clear instructions (README.md) for training and evaluation.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Full training and testing details are in appendix. Full implementations of generative models and classifiers are included in code.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: For results the main experiment (LAMA- AKL benchmark, Llama2-7B baselines), we present error ranges at #2a derived from the number of trials (5), on Ap- pendix A.4.4.
8. Experiments Compute Resources

---

For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We include the name of the GPU we used in addition to the time of execution.
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics htttpes//eerrrpssccc/pbllllleeee
lines
Answer: [Yes]
Justification: We do not work with human participants, and all relevant datasets have been checked for privacy compliance prior to experiments and submission. We do not release any model that could be considered high-risk, and we offer a discussion of broader societal impacts in our discussion section.
Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss the potential positive and negative impacts this work $5
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our contribution does not include new datasets or pre-trained models that pose risk of misuse.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Code that derive from earlier work is properly licensed and referenced.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide anonymized code for our quantitative experiments alongside clear instructions for training and evaluation.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: No human subjects crowdsourcing involved in this research.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human
Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human subjects were involved in this research.
25


Reviews:
[{"review_id": "WvwVoMuQay", "reviewer": "Reviewer_CkYa", "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.", "questions": "Can the authors provide more details on the potential applicability of TAALM to other types of continual learning tasks outside of language models?\nHow does TAALM perform with significantly smaller models or in environments with limited computational resources?\nCould the authors include more detailed ablation studies to better understand the contribution of each component of TAALM to the overall performance?", "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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Originality: The paper introduces a novel approach to CKL by using meta-learning to predict token importance, which is a significant departure from traditional methods.\nQuality: The technical claims are well-supported by both theoretical justifications and empirical evidence. The experiments are comprehensive and rigorously conducted.\nClarity: The paper is well-structured and clearly written, with detailed explanations of the methodology and results.\nSignificance: The contributions are substantial, offering both a new method (TAALM) and a new benchmark (LAMA-CKL) that advance the state-of-the-art in CKL research.", "weaknesses": "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.\nComputational 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.\nAblation 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We truly appreciate your deep understanding of our paper. We address the key issue raised by the reviewer in the comments below:\n***\n**[W1+Q1] Can the authors provide more details on the potential applicability of TAALM to other types of continual learning tasks outside of language models?**\n\nThank you for the constructive feedback. The applicability of TAALM to others is also our main concern. However, we first wish to explain that continual \"knowledge\" learning (CKL) is a critically important issue, while it has been relatively understudied.\n\n**1)TA handles the unique problem of LM**\n\nThe CKL problem is mainly related to LM because possession of knowledge is a unique characteristic of LM. This issue is distinct from traditional continual learning (CL), which primarily focuses on learning new \"tasks\" rather than new \"knowledge\". When we ask well-known large LMs (e.g., chatGPT) \"Who is the current president?\", most reply that, they cannot answer as their knowledge was updated last in 2023 or any previous years. This issue arises because the models lack the capability for CKL. \n\nLM is also different from other domains as it learns sequences. This property yields unique inefficiencies, resulting in extra forgetting. TAALM is specifically designed to tackle this problem. Despite the importance and uniqueness of this CKL issue, research has been conducted mainly by importing general CL approaches from fields outside of LM. That is why this specified study is valuable.\n\n**2)Applicable for other sequence learning: reinforcement learning**\n\nSince our methodology is specialized for sequence learning, we expect it to be applicable for reinforcement learning (RL), which also involves learning of sequence (i.e., Markov Decision Process). Specifically, it seems suitable for addressing the credit assigning problem of sparse reward tasks.\n\nIn the field of RL, one of the major challenges is to learn tasks with sparse rewards. Sparse reward tasks are those where rewards are given only at the end of very long trajectories. Traditional methods like temporal difference learning have a problem with the short observation window, which is considered not suitable to handle this problem. [1]\n\nThe common solution is imitation learning, which is memorizing an entire successful demonstration trajectory without considering rewards. However, it forces a model to absorb all inefficient movements included in the demonstration without self-refining. To solve this, it is necessary to elucidate the impact of each action step (i.e., credit assignment). [2]\n\nOur Train-Attention (TA) is suitable for addressing this credit assignment. TA predicts the importance of each token based on its impact on the final performance, which is very similar to predicting the importance of each action step of a long trajectory. \n***\n**[W2+Q2] How does TAALM perform with significantly smaller models or in environments with limited computational resources?**\n\nThank you for the insightful question. This issue is a primary concern for us, and we aim to address it in the subsequent study.\n\nA promising approach is utilizing a 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.\n\nSince 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.\n\nOn 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. We will make sure to include this in the camera-ready version.\n\n- **Baselines with large generation model (Llama2 7B)**\n| | Parameter size of TA | Top Acc |  Epoch | NF Acc | Total Knowledge | \n|---|---|---|---|---|---|\n| Finetune | NA | 0.1150 | 16 | 0.8174 | 0.9324 |\n| TA (Llama) | 1.1B | 0.4290 | 4 | 0.8983 | 1.3273 |\n| TA (BERT) | 108M | 0.3210 | 6 | 0.9388 | 1.2598 |\n\n- **Baselines with small generation model (Tinyllama 1B)**\n| | Parameter size of TA | Top Acc | Epoch | NF Acc | Total Knowledge |\n|---|---|---|---|---|---|\n| Finetune | NA | 0.0700 | 29 | 0.7693 | 0.8393 |\n| TA (Llama) | 1.1B | 0.3260 | 4 | 0.9078 | 1.2338 |\n| TA (BERT) | 108M | 0.2440 | 9 | 0.9267 | 1.1707 | \n***\n**[W3] 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.**\n\nBased 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. Detailed descriptions of the components and experimental results can be found in the **global rebuttal**.\n***\n\n_**References**_ \n[1] Sqil: Imitation learning via reinforcement learning with sparse rewards \n[2] Learning implicit credit assignment for cooperative multi-agent reinforcement learning", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "knppBdpt7g", "reviewer": "Reviewer_6R4i", "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.", "questions": "- What will be the influence of dropping the lightweight token sequences in training? If only important token sequences matter in learning new knowledge.\n- Can authors offer more theoretical or experimental explanations for the effectiveness of the trained attention?\n- 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? \n- 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": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper introduces Train-Attention Augmented Language Model (TAALM) to reweight token sequences in tuning LLM to avoid catastrophic forgetting. \n- It introduces a new benchmark, LAMA-CKL, for clearer learning-retention trade-off assessment. \n- TAALM demonstrates state-of-the-art performance, compatibility with existing methods, and computational efficiency, advancing the field of CKL.", "weaknesses": "- 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\".\n- This paper lacks of explanation of why it could maintain previous knowledge without forgetting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "First of all, we appreciate your constructive feedback on our work. However, it seems there might be some misunderstandings about our work. Most of the queries raised are addressed in Section 2; we recommend reexamining this section for clarification. Before addressing the issues raised by the reviewer, we would like to provide a summary of our work.\n\n**1)Mature LM learns not only necessary tokens**\n\nOur research aims to address the inefficiencies of the traditional fine-tuning procedure of LMs, thus enhancing its continual knowledge learning (CKL) capacity. What are the inefficiencies in the fine-tuning procedure? Even though the model only needs a part of the information in a sentence, it has to learn all the tokens in the whole sentence.\n\nAs shown in Figure 1 of our paper, let's assume a language model (LM) learns the sentence \"the president of the US is Biden\". If this LM were a scratch model that hasn't learned anything yet, it would need to learn the information in every token sequence, as all sequences contain at least information on grammar rules. However, if this model were a mature model already pretrained and finetuned, it would not need all the information. If this model were trained in 2020 and incorrectly believes that the current president is Trump, then in this sentence, the only token that carries important information would be \"Biden.\"\n\n**2)Restrict unnecessary parameter changes to avoid forgetting**\n\nWhat if we force this mature model to still learn every token in the sentence? Learning means adjusting parameters, and adjusting parameters means forgetting other knowledge which is contained in those parameters. On the contrary, what if we limit the learning scope to only the essential parts and restrict the model from learning the rest? By minimizing parameter adjusting, the model could keep its other knowledge from forgetting. (**W2, Q3**)\n\n**3)Train-Attention (TA) decides what is necessary**\n\nThe role of token importance (meta-weight) is to determine which tokens are important. In the example above, a high weight would be assigned to \"Biden,\" and no weight would be assigned to the other tokens. To determine the “importance” of each token, we propose “usefulness” as a new criterion. Just as a person might remember the name \"Biden\" to avoid being foolish in future conversations, a model learns specific knowledge only to perform well in future tasks. Therefore, we use meta-learning to identify which tokens contain information that is useful for future tasks. (**Q2, Q3**)\n\nWe have demonstrated this sufficiently through experiments. First, through an experiment with the oracle label (Figure 8), we showed that focusing solely on necessary information improves learning speed, capacity, and retention. Additionally, our main experiment (Figure 7, Table 1) proved that the weights predicted by our TA are well-predicted and perform similarly to the oracle. (**Q2**)\n***\n**[W1] 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 \".**\n\nThank you for recommending an excellent paper. However, it is essentially different from our research.\n\nFirst, MetaWeighting is a method of assigning weights to each task when conducting multi-task learning. This work can be categorized as “hard mining” (e.g., hard example mining, hard task mining), and the motivation of this work is to refine and automate this hard mining process using meta-learning. Methodology: MetaWeighting adjusts the learning rate of each task, and within the same task, all data are learned with a uniform weight.\n\nOn the other hand, our work is intended to improve the model’s CKL capacity by addressing the inefficiencies in the standard finetuning procedure, which is caused by uniform weighting through all tokens. Our contributions are 1) The observation that uniform weighting is particularly detrimental to CKL, 2) suggesting a novel view that redefines token importance as usefulness, and 3) this view enables meta-learning to be applied to this problem. These are the novelty of our work and are not similar to MetaWeighting. Methodology: Weights are applied to the tokens of the data being learned, and are optimized to enhance the performance of individual tasks. This is different from meta-weight where weights are applied to the task itself.\n\nAs described, our work and meta-weight research only share the keywords \"task\" and \"meta-learning\". Beyond these, the motivation, implementation pipeline, and objectives are all different.\n\n***\n\n**[W2] This paper lacks of explanation of why it could maintain previous knowledge without forgetting. [Q2] Can authors offer more theoretical or experimental explanations for the effectiveness of the trained attention? [Q3] 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?**\n\nWe address these issues in the prior part of our response. And each is also explained in our paper, mainly on Section 2.\n\n***\n**[Q1] What will be the influence of dropping the lightweight token sequences in training? If only important token sequences matter in learning new knowledge.  [Q4] 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.**\n\nWe really appreciate you for the insightful suggestion. Investigating this issue has led us to a deeper understanding of our own work. \nBased on the token importance predicted by the TA, various design choices are possible, including your suggestion. We explore and compare the effectiveness of these variations. This study enhances understanding of how generative LM interacts with token weights when learning data. Detailed descriptions of the components and experimental results can be found in the **global rebuttal**.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4vuJJN077b", "reviewer": "Reviewer_Az1R", "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.", "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": "The discussion of limitations is included in the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Consider the expected functionality of tokens in solving related tasks as token importance is novel and well-motivated.\n2. The proposed meta-learning framework is simple yet effective and can be complementary to other methods.\n3. A new benchmark for evaluating continual knowledge learning is established.", "weaknesses": "1. Naming token importance as \"usefulness\" is too broad and does not accurately reflect the motivation of the paper.\n2. Meta learning is a well-known method, the description of the algorithm details in Section 3.2 is however confusing and unclear.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**[W1] Naming token importance as \"usefulness\" is too broad and does not accurately reflect the motivation of the paper.**\n\nThank you for the constructive feedback on the clarity of expression. We consider  \"task-utility\" as it more clearly implies usefulness in relation to the task. \n***\n**[W2] Meta learning is a well-known method, the description of the algorithm details in Section 3.2 is however confusing and unclear.**\n\nFor your feedback, we briefly restate using the frame of legacy meta-learning work such as MAML [1], meta-SGD [2] as follows:\n\nMeta-learning consists of an inner loop that trains the model and an outer loop that trains the meta-learner. The true goal is to train the meta-learner. In our work, the meta-learner is called Train-Attention (denoted ϕ), and the model is a generative LLM (denoted θ). \n\n**Inner loop**: θ learns from data (D), updated into θ′. When learning, θ utilizes the token importance (W_ϕ), which is predicted by ϕ, in a target-weighted manner. See Eq.(4)\n\n**Outer loop**: ϕ learn to generate W_ϕ to maximize task (T_D) performance of θ′. See Eq.(5)\n\nEach time the outer loop is executed, ϕ is updated, and θ is repeatedly reset to the initial point.\n\nCompared to [1][2], the meta-learner in MAML is the initial parameter of that model, and the meta-learner in meta-SGD is a hyperparameter (e.g., learning rate). In contrast, our meta-learner is a separate model that predicts the hyperparameter.\n\nIn the camera-ready version, we will adopt this frame to explain more clearly.\n***\n**[W3] The proposed method seems to only be applicable to the experimental setup described in the paper. This setup requires pairs of D and T_D, where T_D represents a task that can be solved using the information contained in D. This setup is not general enough and is limited to very specific scenarios** \n\nWe appreciate the reviewer for highlighting this fundamental question. We assume that our method is applicable on the scenario where LM should continually update its knowledge. And a pair of D and T_D represent quite a general situation, more than it looks.\n\nIt is easy to think that information exists in isolation from practical use, but its value of existence comes up when it finds somewhere to be used. In other words, every data (D) potentially has its task (T_D), and how to find them and link them as a pair is up to future research. We introduce practical examples below.\n\n**1)The common retrieval augmented QA scenario is applicable**\n\nAny retrieval augmented question answering (QA) scenarios are suitable for applying TAALM. Currently, the main solution for this scenario is to retrieve the related document and append it to the prompt while generating answers. However, if we regard the retrieved document as D and the question-answer pair as T_D, it will build an TAALM that is capable of continual knowledge updates through retrieving. Since most current uses of LMs fall into this retrieval-augmented QA scenario, we believe our methodology has significant applicability.\n\n**2)Multi-session dialogue scenario is also applicable**\n\nFurthermore, depending on how we interpret D and T_D, the application of TAALM can become even broader and more general. Consider a scenario involving multiple conversation sessions with our friends. For instance, if a friend says, \"I'm moving to Washington\", in the next meeting, we should remember this move to have a good conversation. In this case, the earlier conversation can be interpreted as D, and the utterance generation task of the subsequent conversation as T_D.\n\nIn practice, we successfully trained TAALM using this approach on the Multi-Session-Chat (MSC) [3] dataset, treating one session as D and the next as T_D. This training approach enhances the model’s dialogue understanding performance in the subsequent sessions compared to standard fine-tuning, which we plan to describe in the following research. We provide a token importance heat map, which is generated by the TA fitted to the MSC dialogue (Figure 3 in the rebuttal pdf).  And we will introduce this approach in the camera-ready version.\n \n**[Q1] Why does the proposed meta learning approach reach performance peaks in fewer epochs compared to other methods? Is there an explanation for this?**\n\nWe truly appreciate the reviewer's careful observation. We provide an easy and detailed version of explanations here.\n\n**Easy explanation**\n\nAs the TA is optimized to predict token weights that maximize the task performance of LM, the LM reaches the performance peak in fewer epochs with the optimal token weights.\n\n**Detailed explanation**\n\nWe hypothesize that the interference among the gradient vectors occurs in the standard finetuning, resulting in slower learning.\n\nFor instance, assume there is a sentence to learn, and five chunks of information (denote A, B, C, D, E) in that one sentence. Among the five, only A is useful and worth learning, while others are not. Assuming the model size as 4-dimension for simplicity, the gradient vector (necessary parameter change for the model to learn each piece of information) for each are: A as [1,1,0,-1], B as [1,1,0,1], and C as [-1,-1,1,0],  et cetera. In this case, the directions of the first and second elements for A and C cancel each other, and the last element of A and B cancels. These interferences will delay the parameter update to the necessary amount to reach A.\n\nAs seen in this case, the more diverse the information to be learned, the more probable that gradient vectors will cancel each other out. This results in extended training steps required to learn the targeted information A. In contrast, TAALM makes the model selectively learn only targeted information, A, thus avoiding interference from other data and enabling faster learning.\n***\n_**References**_\n[1] Model-agnostic meta-learning for fast adaptation of deep networks\n[2] Meta-sgd: Learning to learn quickly for few-shot learning\n[3] Beyond Goldfish Memory: Long-Term Open-Domain Conversation", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hlg9F3nVyi", "reviewer": "Reviewer_dYuZ", "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.", "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.\n\n2. Could the authors provide more analysis and rationale behind the phenomena that K-adapter keeps its not-to-forget accuracy basically unchanged?", "limitations": "The authors addressed the limitations adequately.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. It is inspiring to use meta-learning technique for adjusting token weight.\n\n2. The proposed benchmark improves on distinguishing plasticity and stability.\n\n3. The experiments are comprehensive and solid.", "weaknesses": "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.\n\n2. 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.\n\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**[W1] The idea of adjusting token weights is not fresh enough and shares some similarity with “RHO-1”. However, training a meta-learner to evaluate token importance is still a good try.**\n\nWe truly appreciate your recognition of our effort, as well as your deep understanding of our paper.\n\nAs you mentioned, in addition to the concurrent work RHO-1, earlier studies such as token dropping [1] have also addressed methods for adjusting token weights, and these works are highly inspiring. However, we introduced our work’s novelty from these methods in Section 2, and the performance comparison with the RHO-1 in Section 4.\n\nFor our best understanding, RHO-1 assumes the loss difference between the current training model and the reference model as a measure of importance. It indicates that lower confidence correlates with higher importance. This shares similar concept with previous work of token dropping [1] and focal loss [2].\n\nHowever, these methods still allow many tokens to be learned that a sufficiently trained mature model does not need, mainly due to its naive definition of “token importance”. Therefore, these techniques have been introduced more as pretraining methods rather than as CKL approaches. \n\nIn contrast, we define token importance as \"usefulness\" rather than confidence, and propose a perspective that the acquisition of knowledge should be optimized for the task. This novel perspective enables meta-learning methods to be applied to token weighting. This is the key distinction between our method and previous token weighting approaches.\n\nOur work's effectiveness is also demonstrated through experiments. The result in Table1 shows that our approach is significantly more optimal than RHO-1 for CKL scenario.\n***\n**[W2] 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.**\n\nFor your suggestion, we include heat-maps of token importance from TA, in the **rebuttal pdf**.\n\nTA 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 1). In passages listing regional relations, TA pinpoints the names of locations (Figure 2). 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).\n\nWe also provide an attention map of TA trained on the multi-session chat [3] (Figure3). 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.\n***\n**[Q1] In Figure 7, I wonder why the to-learn accuracy drops after the fourth epoch. It is kind of counterintuitive to me, as I expect it to grow along with the increasing epoch.**\n\nAs the reviewer observed, in the main experiment (Figure 7), the performance declines after a rapid peak.\n\nWe 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.\n\nThe comparison between TA and the oracle (Figure 8) 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.\n\nThis 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.\n***\n**[Q2] Could the authors provide more analysis and rationale behind the phenomena that K-adapter keeps its not-to-forget accuracy basically unchanged?**\n\nAs depicted in Figure 7, K-Adapter shows a slower increase in to-learn accuracy and a slower decrease in not-to-forget accuracy compared to baselines using QLoRA adapter. This phenomenon appears to be due to structural differences between K-Adapter and QLoRA.\n\nQLoRA, a quantized version of Lora, adds a new parameter, C (the result of multiplying two smaller matrices, A and B), to the original model parameters. As the original parameters can't be recovered after C is added, Lora is essentially the same as modifying the original parameters. Modifying parameters means that the model forgets other knowledge that is previously held by the parameters.\n\nIn contrast, the K-Adapter is a separate transformer layer that takes the hidden states of the original model as inputs. This layer includes features like gates and residual networks, which can choose to let the original outputs pass through without changing them.\n\nTo summarize, while QLoRA modifies the existing model parameters and stores only the changes in the adapter, while K-Adapter does not directly alter the parameters but instead trains an additional independent transformer alongside them. This approach makes K-Adapter slower in learning, but it seems to retain more of the original knowledge.\n***\n_**References**_\n[1] Token dropping for efficient bert pretraining    \t\n[2] Focal loss for dense object detection\n[3] Beyond Goldfish Memory: Long-Term Open-Domain Conversation", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# Global Rebuttal


We first thank all reviewers for their thoughtful feedback on our work. We would like to address a suggestion commonly raised by reviewers, and introduce our progress in significantly reducing the GPU resource required. 
We believe that constructive suggestions from all reviewers, such as the experiment on the token dropping method can significantly enhance the clarity and advancement of our paper.
we will make sure to include the updates in the camera-ready version
***
## **§A. Ablation study**
Based on the token importance predicted by the Train-Attention (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. (suggested by reviewer **6R4i**)

### **Result**:
Known token masking does not yield better results compared to TAALM w/ token-importance 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.

- **Baselines with large generation model (Llama2 7B)**

| | Top Acc | Epoch | NF Acc | Total Knowledge |
|---|---|---|---|---|
| Finetune | 0.1150 | 16 | 0.8174 | 0.9324 |
| TAALM w/ token-importance weight (**ours**) | 0.4290 | 4 | 0.8983 | 1.3273 |
| TAALM w/ known token masking  | 0.3920 | 4 | 0.9075 | 1.2995 |
| TAALM w/ token weight dropping < 0.5 | 0.4100 | 7 | 0.9148 | 1.3248 |
| TAALM w/ token weight dropping < 0.8 | 0.3850 | 4 | 0.9267 | 1.3117 | 

***
## **§B. Reduction of resources through TA on BERT**
Due to the substantial GPU resources required to train the TA, we endeavored 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. We will make sure to include this in the camera-ready version.

- **Baselines with large generation model (Llama2 7B)**
| | Parameter size of TA | Top Acc | Epoch | NF Acc | Total Knowledge | 
|---|---|---|---|---|---|
| Finetune | NA | 0.1150 | 16 | 0.8174 | 0.9324 |
| TA (Llama) | 1.1B | 0.4290 | 4 | 0.8983 | 1.3273 |
| **TA (BERT)** | 108M | 0.3210 | 6 | 0.9388 | 1.2598 |


- **Baselines with small generation model (Tinyllama 1B)**
| | Parameter size of TA | Top Acc | Epoch | NF Acc | Total Knowledge |
|---|---|---|---|---|---|
| Finetune | NA | 0.0700 | 29 | 0.7693 | 0.8393 |
| TA (Llama) | 1.1B | 0.3260 | 4 | 0.9078 | 1.2338 |
| **TA (BERT)** | 108M | 0.2440 | 9 | 0.9267 | 1.1707 |

Meta-review: 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.
Decision: Accept (poster)