[2025-2] 전연주 - Train-Attention: Meta-Learning Where to Focus in Continual Knowledge Learning

논문 링크: https://arxiv.org/abs/2407.16920

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

Conference: NeurIPS 2024 Poster

2407.16920v1.pdf

7.26MB

---

1. 문제의식: CKL에서 “모든 토큰을 똑같이” 학습하는 게 문제다

• LLM을 continual knowledge learning(CKL) 설정에서 업데이트하면
  새로운 지식을 배우는 동안 기존 지식을 많이 잊는 catastrophic forgetting 문제가 생김.
• 기존 방법들(regularization, adapter, rehearsal 등)은
  구조·규제는 다르지만 공통적으로
  “모든 토큰에 동일한 weight를 주고 학습”한다는 비효율을 그대로 물려받음.
• 실제로는 한 문장 안에서도
  진짜 새로운 지식을 담은 토큰은 일부(예: “US 대통령은 Biden이다”에서 “Biden”),
  나머지는 문법/이미 아는 상식이라 굳이 또 업데이트할 필요가 없음.

---

2. 아이디어: 중요한 토큰만 골라 더 세게 학습하자 (Train-Attention + TWL)

논문이 제안하는 핵심 아이디어는 두 가지입니다.

1. Token-Weighted Learning (TWL)
   • 기존 LM의 loss: $$\sum_i \log p(x_i \mid x_{<i})$$
   • 여기서 각 토큰마다 중요도 weight $w_i$​ 를 곱해주는 방식으로 바꿈.
   • 중요한 토큰은 크게 업데이트, 덜 중요한 토큰은 거의 안 건드림.
2. Train-Attention (ϕ)
   • 어떤 토큰이 중요한지 사전에 정해져 있는 것이 아니라,
   • 별도 모델(Train-Attention)이 문서 D를 보고 각 토큰의 weight를 예측함.
   • 이 ϕ를 meta-learning 방식으로 학습:
     • inner loop: ϕ가 예측한 weight로 base LM(θ)을 한 스텝 학습
     • outer loop: 그렇게 업데이트된 θ′가 실제 task $T_D$​ 를 얼마나 잘 푸는지를 보고
       ϕ의 파라미터를 업데이트
   • 결과적으로 ϕ는
     “이 토큰에 집중하면 나중에 task 성능이 가장 많이 좋아지는 weight”를 배우게 됨.

---

3. 새로운 벤치마크: LAMA-CKL

• 기존 TEMPORALWIKI 벤치마크에서는
  CHANGED/UNCHANGED 성능이 같이 오르락내리락해서
  plasticity(새 지식 학습) vs stability(기존 지식 유지)의 trade-off가 명확히 안 보였음.
• 그래서 저자들이 LAMA-CKL이라는 새 벤치마크를 설계:
  • LAMA TREx에서
    • TO-LEARN set: 시간에 따라 바뀌는 relation + 기존 모델이 못 맞추는(Acc=0) triple
    • NOT-TO-FORGET set: 변하지 않는 relation + 이미 잘 맞추는(Acc=1) triple
  • 500개씩 작은 데이터로 30 epoch 동안 반복 학습
    → TO-LEARN accuracy는 올라가고, NOT-TO-FORGET accuracy는 내려가는 패턴이 명확히 보임
    → CKL의 본질적인 trade-off를 잘 드러내는 벤치마크.

---

4. 실험 결과: TAALM은 “많이·빨리·덜 잊으면서” 배운다

• Base LM: Llama2-7B + QLoRA 등 다양한 세팅.
• 비교 대상:
  • Finetune(QLoRA)
  • K-Adapter
  • Mix-review
  • RecAdam
  • RHO-1 (다른 token selection 방법)
  • Oracle weight (object token만 1로 주는 이상적인 upper bound)

LAMA-CKL 결과:

• TAALM(QLoRA)의 TO-LEARN Top Acc = 0.429 (4 epoch만에 도달)
  • 기존 2위(RHO-1=0.141)의 3배 이상
  • Epoch=4라서 다른 방법보다 훨씬 빨리 수렴
• 같은 시점의 NOT-TO-FORGET Acc도 0.898 수준으로 잊어버림이 상대적으로 적음.
• Oracle 대비 Top Acc의 약 78% 수준까지 접근
  → meta-learning으로 예측한 weight가 인간이 직접 지정한 “정답 weight”에 꽤 근접.

TEMPORALWIKI 결과:

• TinyLlama-1B 기준, QLoRA·K-Adapter 베이스 실험 모두에서
  TAALM이 가장 낮은 perplexity(=최고 성능) 기록.
• LAMA-CKL용으로 학습된 TAALM을 그냥 TEMPORALWIKI에 써도 2등 성능 → 일정 수준의 transferability 입증.

또한, TAALM을 기존 기법(K-Adapter, Mix-review, RecAdam)에 붙이면
모든 baseline 성능이 향상되는 시너지 효과도 보임.

---

5. 한계 및 향후 방향

1. Task-specificity
   • Train-Attention은 학습 당시의 task 분포에 맞춰 중요 토큰을 배우기 때문에,
   • 분포가 매우 다른 완전히 새로운 task에 대한 일반화에는 한계가 있을 수 있음.
   • 하지만 일반적인 “지식 획득” task라면 다른 CKL task에도 어느 정도 전이되는 것으로 보임.
2. Data–task pair 필요
   • Train-Attention을 학습하려면 (D, Tᴅ) 쌍이 필요함.
   • 만약 데이터와 task가 따로 있다면:
     • search로 매칭하거나 (lexical search, dense search 등)
     • LLM을 이용해 synthetic task를 생성하는 방법을 제안.
3. Broader impact
   • CKL을 더 효율적으로 만들고, 전체 LLM 재학습 비용을 줄이는 데 기여 가능.
   • 사회적 부정적 영향은 크지 않다고 판단.

---

* perplexity = 언어 모델이 다음 토큰을 얼마나 덜 헷갈리고 잘 예측하는지를 수치화한 지표.