Research map / Continual learning / Continual fine-tuning of LLMs
Train-Attention: Meta-Learning Where to Focus in Continual Knowledge Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TAALM, which meta-learns a Train-Attention model that predicts per-token weights for continual knowledge learning in causal LMs, and introduces the LAMA-ckl benchmark for plasticity-stability trade-off.
Problem
Continual knowledge learning (CKL) in LLMs suffers from catastrophic forgetting, and standard causal LM training applies uniform weights to all tokens, causing unnecessary parameter updates and forgetting. Existing benchmarks also fail to clearly show the learning-retention trade-off.
Approach
Define token importance as the expected downstream usefulness of a token and meta-learn a Train-Attention head (replacing the LM head with a single-output sigmoid-like scalar per position) that predicts token weights for token-weighted learning; the weights are optimized by the post-update loss on a related task, with the base model reset after each meta-step.
Key results
- On LAMA-ckl with Llama2-7B, TAALM (QLoRA) reaches Top Acc 0.4290 at epoch 4 versus 0.1410 for RHO-1 and 0.1150 for standard finetune, with Total Knowledge 1.3273.
- On TemporalWiki with TinyLlama-1B (QLoRA), TAALM achieves lower perplexity than all baselines across 0910, 1011, and 1112 periods (e.g., avg 4.1435 vs 4.507 for RHO-1 on 0910).
- Oracle token weighting reaches Top Acc 0.5470, and TAALM achieves 78.2% of Oracle Top Acc; a BERT-based Train-Attention (108M) reaches Top Acc 0.3210 on a single 24GB GPU.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply Train-Attention token weighting instead of standard QLoRA finetuning | Llama2-7B · LAMA-ckl | Top Acc: 0.1150 -> 0.4290 | Helped |
| Apply Train-Attention token weighting instead of RHO-1 token selection | Llama2-7B · LAMA-ckl | Top Acc: 0.1410 -> 0.4290 | Helped |
| Apply Train-Attention token weighting instead of K-Adapter | Llama2-7B · LAMA-ckl | NF Acc: 0.9980 -> 0.8983 | Hurt |
| Apply Train-Attention token weighting instead of RHO-1 token selection | TinyLlama-1B QLoRA · TWiki-Probes-0910 Avg | perplexity: 4.507 -> 4.1435 | Helped |
| Apply Train-Attention token weighting instead of RHO-1 token selection | TinyLlama-1B QLoRA · TWiki-Probes-1011 Avg | perplexity: 4.3775 -> 4.092 | Helped |
| Apply Train-Attention token weighting instead of RHO-1 token selection | TinyLlama-1B QLoRA · TWiki-Probes-1112 Avg | perplexity: 4.594 -> 4.197 | Helped |
| Combine TAALM with Mix-review rather than use TAALM alone | Llama2-7B · LAMA-ckl | Top Acc: 0.4290 -> 0.4500 | Helped |
| Use Oracle token weights rather than uniform standard finetuning weights | Llama2-7B · LAMA-ckl | Top Acc: 0.1150 -> 0.5470 | Helped |
| Use BERT rather than Llama as the Train-Attention model | Llama2-7B TAALM · LAMA-ckl | Top Acc: 0.4290 -> 0.3210 | Hurt |
| Mask tokens in real time when the predicted importance matches the label (ablation) | Train-Attention token weights · LAMA-ckl | performance | Hurt |
| Drop weights below the top k% token-importance threshold (ablation) | Train-Attention token weights · LAMA-ckl | performance | Hurt |
Surprising findings
- Heuristic adjustments such as masking tokens already predicted correctly or dropping low-importance tokens degrade performance, which the authors attribute to Train-Attention already being near-optimal.
- Accuracy on the to-learn set declines after an early peak (around epoch 4), which the authors attribute to overfitting on false-target tokens.
Limitations
- Train-Attention is task-specific and requires paired data-task examples, which may not be available in general settings.
- Accuracy declines after peak performance, suggesting overfitting to false token targets.
- Evaluation is limited to LAMA-ckl and TemporalWiki; several baselines are not directly compared due to unavailable benchmarks.
Experimental setup
Datasets: LAMA (T-REx) TemporalWiki Wikipedia (BERT pretraining)
Benchmarks: LAMA-ckl TemporalWiki
Metrics: Top Acc Epoch NF Acc Total Knowledge perplexity on Changed/Unchanged sets
Baselines: Finetune (QLoRA) K-Adapter Mix-review (QLoRA) RecAdam (QLoRA) RHO-1 (QLoRA) Oracle token weight
Builds on: QLoRA K-Adapter RHO-1 TemporalWiki LAMA
Closest papers by meaning
- Learn more, but bother less: parameter efficient continual learning (NeurIPS 2024 · similarity 0.90)
- Scalable Language Model with Generalized Continual Learning (ICLR 2024 · similarity 0.90)
- Spurious Forgetting in Continual Learning of Language Models (ICLR 2025 · similarity 0.89)
- CMT: A Memory Compression Method for Continual Knowledge Learning of Large Language Models (AAAI 2025 · similarity 0.88)
- CoIN: A Benchmark of Continual Instruction Tuning for Multimodel Large Language Models (NeurIPS 2024 · similarity 0.88)
- From RAG to Memory: Non-Parametric Continual Learning for Large Language Models (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.