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

TL;DR

The paper introduces TAALM, a meta-learning framework that predicts token-level usefulness to weight CKL updates, addressing catastrophic forgetting in LLMs. By training a Train-Attention head to assign per-token weights, the approach enables targeted updates and reduced forgetting, validated on the LAMA-ckl and TemporalWiki benchmarks. TAALM demonstrates state-of-the-art CKL performance, synergizes with prior CKL methods, and remains effective with smaller models and resources. The proposed LAMA-ckl benchmark makes the learning–retention trade-off more explicit, guiding future CKL research and practical deployments.

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, \textsc{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.

Figures (18)

• Figure 1: (a) Learning of Causal LM: The document is decomposed into multiple token sequences $s_i\doteq x_i|x_{<i}$, which aligns with different importance, but uniformly weighted. (b) Train-Attention: Our proposed Train-Attention learns to predict weights that approximate importance, to enable targeted continual knowledge updates through label-free meta-learning method.
• Figure 2: (a) depicts the architecture of Train-Attention, which shares the structure of causal LM, while the decoder layer (LM head) of causal LM is replaced from a linear layer of [$hidden\ size \times vocab\ size$] dimension to [$hidden\ size \times 1$] dimension, which is TA (Train-Attention) head. (b) depicts the TAALM, where the Train-Attention ($\phi$) is augmented to the base model ($\theta$).
• Figure 3: Optimal $W$ leads $\theta$ closer to $\theta^*$.
• Figure 4: One step update of $\phi$.
• Figure 5: Evaluation procedure of the LAMA-ckl benchmark.