TokMem: One-Token Procedural Memory for Large Language Models

이 논문은 프롬프트의 반복 처리 과부하와 모듈화 한계를 해결하기 위해, 재사용 가능한 작업 절차를 단일 학습 가능 메모리 토큰으로 컴파일하여 LLM 의 추론을 제어하고 새로운 절차를 기존 모델에 간섭 없이 지속적으로 추가할 수 있는 'TokMem' 프레임워크를 제안합니다.

핵심

🧠 TokMem: 거대 언어 모델의 '한 글자'로 된 만능 지시서

이 논문은 TokMem이라는 새로운 기술을 소개합니다. 쉽게 말해, 거대 언어 모델 (LLM) 이 새로운 일을 배울 때, 긴 설명서 (프롬프트) 를 계속 읽게 하는 대신, 작은 '기억 토큰' 하나만 기억하게 하여 일을 시키는 방법입니다.

이해를 돕기 위해 일상적인 비유를 들어 설명해 드리겠습니다.

---

1. 문제점: "긴 설명서"의 함정

지금까지 AI 에게 새로운 일을 시키려면, 매번 긴 설명서 (프롬프트) 를 붙여주었습니다.

• 비유: 요리사를 시켜 "오늘 저녁 메뉴를 만들어줘"라고 할 때, 매번 수백 페이지 분량의 요리책을 그 요리사 책상 위에 펼쳐놓고 "이거 봐, 이거 보고 해"라고 하는 것과 같습니다.
• 단점:
  • 책상 (메모리) 이 금방 꽉 차서 다른 재료를 넣을 공간이 없습니다.
  • 요리사가 책장을 넘기는 데 시간이 걸려 요리가 늦어집니다 (계산 비용 증가).
  • 같은 레시피를 매번 다시 읽어야 하므로 비효율적입니다.

2. 해결책: TokMem (한 글자 지시서)

TokMem 은 이 긴 설명서를 작은 '기억 토큰' 하나로 압축합니다.

• 비유: 요리사에게 수백 페이지의 책을 주는 대신, 작은 스티커 하나를 붙여줍니다.
  • 스티커에 "건강식"이라고 적혀 있다면, 요리사는 그 스티커만 보고도 "아, 오늘 저녁은 건강식 레시피대로 하면 되겠다!"라고 바로 이해하고 실행합니다.
  • 이 스티커는 실제 텍스트가 아니라, AI 가 이해하는 특별한 '기호'입니다.

3. 어떻게 작동할까요?

🏗️ 1 단계: 새로운 기술 배우기 (학습)

• AI 는 새로운 작업 (예: "건강한 저녁 메뉴 만들기") 을 배울 때, 긴 텍스트 대신 그 작업에 해당하는 **새로운 스티커 (기억 토큰)**를 하나 만듭니다.
• 이 스티커는 AI 의 두뇌 (기존 모델) 를 건드리지 않고, 별도의 기억 은행에만 저장됩니다.
• 장점: 기존 지식을 망가뜨리지 않고 새로운 기술을 계속 추가할 수 있습니다. (기억 은행에 스티커를 계속 붙여나가는 것)

⚡ 2 단계: 일 시키기 (추론)

• 사용자가 "건강한 저녁 메뉴 만들어줘"라고 요청하면, AI 는 기억 은행을 뒤져서 "건강식" 스티커를 찾아냅니다.
• 그 스티커 하나만 붙이면, AI 는 긴 설명 없이도 바로 그 작업을 수행합니다.
• 복잡한 작업의 경우: 여러 스티커를 순서대로 붙일 수 있습니다.
  • 예: [재료 찾기 스티커] → [영양 계산 스티커] → [메뉴 정리 스티커]
  • 마치 요리사가 "재료 먼저 구하고, 칼로리 계산하고, 메뉴 정리해"라는 작업 지시 카드를 한 장씩 넘기며 일하는 것과 같습니다.

4. 왜 이 기술이 특별한가요?

1. 공간 절약 (가방 정리): 긴 설명서 대신 스티커 하나만 쓰므로, AI 의 책상 (메모리) 이 항상 비어있어 더 많은 일을 처리할 수 있습니다.
2. 빠른 속도: 긴 책을 읽을 필요가 없으므로, AI 가 답변을 내는 속도가 훨씬 빨라집니다.
3. 잊어버리지 않음 (기억력): 새로운 스티커를 붙여도 기존 스티커들이 지워지지 않습니다. 마치 새로운 요리 레시피를 배워도 예전에 배운 레시피가 사라지지 않는 것처럼, AI 는 영원히 잊지 않고 계속 배울 수 있습니다.
4. 효율성: AI 의 두뇌 (모델 파라미터) 를 수정하지 않고, 오직 작은 스티커들만 학습시키면 되므로 비용이 매우 저렴합니다.

🌟 요약

TokMem은 거대 언어 모델에게 **"긴 설명서" 대신 "작은 기억 카드 (스티커)"**를 주어, 복잡한 일을 빠르고 정확하게, 그리고 오래 기억하게 만드는 혁신적인 방법입니다.

이제 AI 는 두꺼운 매뉴얼을 들고 헤매는 신입 사원이 아니라, 작은 메모 한 장만 보고도 모든 일을 척척 해내는 베테랑이 된 셈입니다!

기술 요약

TokMem: 대규모 언어 모델을 위한 단일 토큰 절차적 메모리 (One-Token Procedural Memory) 기술 요약

본 논문은 ICLR 2026 에 발표된 **"TokMem: One-Token Procedural Memory for Large Language Models"**으로, 대규모 언어 모델 (LLM) 의 프롬프트 효율성과 지속적인 학습 (Continual Learning) 문제를 해결하기 위해 제안된 새로운 프레임워크입니다.

1. 문제 정의 (Problem)

기존 LLM 은 주로 프롬프트 (Prompt) 를 통해 행동을 제어합니다. 그러나 이 방식에는 다음과 같은 근본적인 한계가 존재합니다:

• 비효율성 및 확장성 부족: 새로운 쿼리마다 긴 프롬프트를 재처리해야 하며, 다양한 작업에 대해 모듈식으로 재사용하기 어렵습니다.
• 계산 비용 증가: 긴 컨텍스트는 자기 주의 (Self-attention) 메커니즘의 이차적 (Quadratic) 복잡도로 인해 추론 비용을 급격히 증가시킵니다.
• 컨텍스트 창 제한: 긴 프롬프트는 입력/출력에 사용 가능한 컨텍스트 창을 축소시켜 정보 손실을 초래합니다.
• 지속적 학습의 어려움: 기존 지식 (파라미터) 을 유지하면서 새로운 작업을 추가할 때, 기존 방식은 종종 '재학습 (Re-training)'이나 '파라미터 간섭' 문제를 일으켜 기존 성능이 저하될 수 있습니다.

기존의 검색 기반 메모리 (RAG 등) 는 텍스트를 검색하여 다시 읽는 방식으로, 여전히 컨텍스트 창을 차지하고 텍스트 해석에 따른 오버헤드가 존재합니다.

2. 방법론 (Methodology)

저자들은 TokMem을 제안합니다. 이는 각 재사용 가능한 작업 절차 (Procedure) 를 **학습 가능한 단일 메모리 토큰 (Trainable Memory Token)**으로 컴파일하는 모듈형 프레임워크입니다.

핵심 메커니즘

1. 단일 토큰 인코딩:

   • 각 작업 절차는 별도의 가중치 벡터 $m_i$로 인코딩되어 어휘집 (Vocabulary) 에 특수 토큰으로 추가됩니다.
   • 이 토큰은 두 가지 역할을 수행합니다:
     1. 인덱싱 (Indexing): 특정 절차를 호출하는 주소 역할.
     2. 조향 (Steering): 생성 과정을 유도하는 제어 신호 역할.
   • 이를 통해 복잡한 절차도 하나의 토큰으로 표현되어 오버헤드가 상수 크기 (Constant-size) 로 유지됩니다.

2. 동결된 백본 (Frozen Backbone):

   • LLM 의 기본 백본 파라미터는 고정 (Frozen) 된 채로 유지됩니다.
   • 오직 메모리 토큰의 임베딩 벡터만 학습됩니다. 이는 새로운 절차를 추가할 때 기존 지식을 손상시키지 않고 모듈식으로 확장할 수 있게 합니다.

3. 추론 및 체이닝 (Inference & Chaining):

   • 단일 절차: 쿼리가 들어오면 모델은 최종 은닉 상태에서 가장 확률이 높은 메모리 토큰을 예측하여 해당 절차를 호출합니다.
   • 복합 절차 (Compositional): 여러 단계가 필요한 작업의 경우, 모델은 각 단계가 완료된 후 다음 단계에 필요한 메모리 토큰을 순차적으로 예측하여 체이닝합니다 (예: 파싱 -> 검색 -> 포맷팅).

4. 안정화 전략 (Renormalization):

   • 지속적인 학습 과정에서 새로 추가된 메모리 토큰의 노름 (Norm) 이 커져 기존 토큰을 압도하는 '노름 팽창 (Norm Inflation)' 문제를 방지하기 위해, 새로운 임베딩을 기존 메모리 뱅크의 평균 스케일에 맞춰 재조정 (Renormalization) 하는 경량화 기법을 도입했습니다.

3. 주요 기여 (Key Contributions)

• 절차적 지식의 컴팩트한 표현: 복잡한 작업 절차를 텍스트가 아닌 단일 학습 가능 토큰으로 압축하여, 컨텍스트 창 소모 없이 효율적인 절차적 제어를 가능하게 함.
• 파라미터 효율성 및 지속적 학습: 백본 모델을 수정하지 않고 오직 토큰 임베딩만 학습함으로써, 새로운 작업을 추가해도 기존 작업의 성능 저하 (Catastrophic Forgetting) 를 방지하고 지속적 학습을 지원함.
• 모듈형 구성 (Composability): 단일 토큰을 체이닝하여 복잡한 다단계 작업 (예: 함수 호출 체이닝) 을 유연하게 수행할 수 있는 아키텍처 제안.

4. 실험 결과 (Results)

저자들은 **Super-Natural Instructions (SNI)**의 1,000 개 작업 (원자적 회상) 과 APIGen 데이터셋의 함수 호출 (복합적 회상) 을 통해 TokMem 을 평가했습니다.

주요 성과

1. 원자적 회상 (Atomic Recall):

   • 1,000 개의 작업이 누적되어도 TokMem 은 RAG(검색 기반) 나 기존 파인튜닝 (Fine-tuning) 보다 높은 성능을 유지했습니다.
   • 특히 RAG 는 작업 수가 증가함에 따라 성능이 급격히 저하되는 반면, TokMem 은 일관된 성능을 보였습니다.
   • 라우팅 정확도: 1,000 개 작업 중에서도 TokMem 은 94% 이상의 정확한 메모리 토큰 선택을 보여주었으며, 이는 RAG 의 검색기 (Sentence-BERT) 보다 훨씬 우수했습니다.

2. 복합적 회상 (Compositional Recall):

   • 여러 도구를 조합하여 문제를 해결하는 함수 호출 태스크에서 TokMem 은 ICL(문맥 학습) 과 RAG 를 압도했습니다.
   • 파라미터 효율성: LoRA 와 같은 파라미터 효율적 파인튜닝과 유사하거나 더 높은 성능을 내면서, 학습 가능한 파라미터 수는 10 배 이상 적게 사용했습니다 (예: Llama 8B 기준 LoRA 는 3.41M 파라미터 필요, TokMem 은 0.2M).
   • 제너럴라이제이션: 훈련 시 보지 못한 더 긴 도구 호출 체인 (Zero-shot) 에 대해서도 TokMem 은 파인튜닝보다 우수한 일반화 능력을 보였습니다.

3. 데이터 효율성:

   • 소량의 학습 데이터 (Few-shot) 환경에서도 TokMem 은 파인튜닝보다 빠른 수렴과 높은 성능을 보여주었습니다.

5. 의의 및 결론 (Significance)

TokMem 은 LLM 의 메모리 아키텍처에 있어 **선언적 메모리 (Declarative Memory, 텍스트 기반)**에서 **절차적 메모리 (Procedural Memory, 토큰 기반)**로의 패러다임 전환을 제시합니다.

• 효율성: 긴 프롬프트 없이도 복잡한 작업을 수행할 수 있어 추론 비용과 지연 시간을 크게 줄입니다.
• 확장성: 새로운 작업을 추가할 때마다 모델을 재학습하거나 컨텍스트를 늘릴 필요 없이, 단순히 새로운 토큰을 추가하는 것만으로 시스템을 확장할 수 있습니다.
• 실용성: 지속적인 학습이 필요한 실제 애플리케이션 (예: 개인화 비서, 동적 도구 호출 시스템) 에서 LLM 의 안정성과 적응력을 동시에 확보할 수 있는 강력한 솔루션을 제공합니다.

결론적으로, TokMem 은 LLM 이 인간과 유사하게 '기술 (Skill)'을 토큰 단위로 학습하고 재구성할 수 있게 함으로써, 차세대 효율적이고 적응적인 AI 에이전트 개발의 중요한 기반이 될 것으로 기대됩니다.